Skip to content

← JAX Performance & Distributed step 2 of 25

Easy Primitives

jit with static_argnames

Jump to the task

Why this matters

A static argument has a value JAX can use while tracing a function. For tiling, this lets the repeat count determine the output length. You can select that argument by position with static_argnums or by parameter name with static_argnames.

Naming the parameter makes the intent easier to read. It also avoids updating a numeric index when you reorder parameters, though the calls must still match the new function signature.

Worked mini-example

import jax
import jax.numpy as jnp

def tile_vector(x, n):
    return jnp.tile(x, n)

tile = jax.jit(tile_vector, static_argnames=("n",))
x = jnp.array([1.0, 2.0])
a = tile(x, n=3)  # [1, 2, 1, 2, 1, 2]
b = tile(x, 2)    # [1, 2, 1, 2]

Both keyword and positional calls work here. When only static_argnames is supplied, JAX inspects the function signature to find the corresponding positional arguments.

Each new count may need another compiled version. Returning to a previous count can reuse its cached version if the other compilation inputs, such as array shape and dtype, still match.

Common pitfalls

  • Use the exact, case-sensitive parameter name from the function passed to jax.jit. If that helper calls its count reps, use static_argnames=("reps",).
  • A static count should be a Python integer, not a scalar JAX array: static values must be immutable and hashable.
  • Define and wrap the helper once. Creating a new inner function on every call prevents reuse based on that function's identity.
  • Convert the fixture's float count before calling the helper. An outer JIT also needs its count marked static; the inner declaration does not make an outer dynamic tracer concrete.

Problem

Implement jit_with_named_static(x, n). Repeat the whole one-dimensional array x exactly n times using jnp.tile in a helper wrapped with jax.jit(..., static_argnames=("n",)).

The runner supplies a nonnegative, integer-valued float for n. Convert it to a Python integer before passing it to the helper. Return shape (x.shape[0] * int(n),). Zero repeats or an empty input give an empty array.

For example, x = [5, 6] and n = 2.0 give [5, 6, 5, 6], not [5, 5, 6, 6].

The diagram marks the second parameter as static. In this exercise, selecting it by the name n has the same effect.

Loading visualization…