Skip to content

← JAX Foundations step 3 of 25

Medium End-to-End

Jit with static_argnames

Jump to the task

Why this matters

Repeating [7, 8] twice gives [7, 8, 7, 8]. The repeat count determines the output length. JAX needs that length while compiling, before it runs the operation on actual array values.

A normal JIT argument does not provide its value during tracing. A static argument does. Use static_argnames=("n",) to make the parameter named n available at trace time.

Worked mini-example

import jax
import jax.numpy as jnp

def repeat_vector(x, n):
    return jnp.tile(x, int(n))

repeat_jit = jax.jit(repeat_vector, static_argnames=("n",))
x = jnp.array([7.0, 8.0])
a = repeat_jit(x, 2.0)  # [7, 8, 7, 8]
b = repeat_jit(x, 3.0)  # [7, 8, 7, 8, 7, 8]
c = repeat_jit(x, 2.0)  # can reuse the cached version for n=2.0

Here n stays a Python float during tracing, and int(n) converts it to the integer that tile needs. Declaring an argument static does not change its type.

Without the static declaration, int(n) would try to read the value of a dynamic tracer and fail with ConcretizationTypeError.

Common pitfalls

  • Static values must be immutable and hashable. Python numbers and strings work; lists and JAX arrays do not. A tuple works only if its contents are also hashable.
  • Converting with int works on a concrete count. It cannot make a dynamic JIT argument static.
  • A new static value may need another compiled version. Prefer a small set of counts if you want compilation costs to stay modest.
  • static_argnames refers to the wrapped function's parameter name. For an ordinary function like this one, either positional or keyword calls work.
  • jnp.tile repeats the whole vector. jnp.repeat repeats individual elements and gives a different order.

Problem

Implement jit_with_static_n(x, n) to tile the one-dimensional array x exactly n times. The count is nonnegative and integer-valued; it may arrive as a Python float. Convert it with int(n) and return shape (x.shape[0] * int(n),). Zero repeats or an empty input produce an empty array.

The function itself does not need to call jax.jit. The runner checks its output directly. To try compilation yourself, wrap your implementation with jax.jit(jit_with_static_n, static_argnames=("n",)). Keep this wrapped function for later calls.

The diagram calls its static count k; it plays the same role as n here. Returning to the first count can reuse its version.

Loading visualization…