Skip to content

← JAX Performance & Distributed step 1 of 25

Medium Primitives

jit with static_argnums

Jump to the task

Why this matters

Tiling repeats a whole array. For x = [1, 2], tiling three times gives [1, 2, 1, 2, 1, 2]. The repeat count determines the output length, so JAX needs to know it when compiling this operation.

By default, JIT traces array and scalar inputs without using their actual values. Marking the repeat count as static makes its value available during tracing. static_argnums=(1,) selects the second argument; positions start at zero.

Worked mini-example

import jax
import jax.numpy as jnp

tile = jax.jit(jnp.tile, static_argnums=(1,))
x = jnp.array([1.0, 2.0])
a = tile(x, 2)       # [1, 2, 1, 2]
b = tile(x, 3)       # [1, 2, 1, 2, 1, 2]
c = tile(x + 10, 2)  # [11, 12, 11, 12]

The first two calls need versions for two different repeat counts. The third can reuse the version for two repeats: the array values changed, but its shape and dtype stayed the same. This assumes the cached version is still available.

Common pitfalls

  • jnp.tile repeats the whole vector. jnp.repeat(x, n) repeats each element, which gives a different order.
  • Static values must be immutable and hashable. Use a Python integer for the count, not a scalar JAX array.
  • Many distinct counts can mean substantial compilation time and cache memory. Static arguments work best with a small set of values.
  • Check the position again if you reorder a helper's parameters. static_argnames can identify the parameter by name instead.
  • Convert the fixture's float count to an integer before calling the jitted helper. int(...) cannot turn a dynamic JIT tracer into a known value. If you also JIT the outer function, mark its count static too.

Problem

Implement jit_with_static_int(x, n_repeats). Tile the one-dimensional array x using jnp.tile wrapped with jax.jit(..., static_argnums=(1,)).

The test runner supplies a nonnegative, integer-valued float for n_repeats, such as 3.0. Convert it to a Python integer before passing it to the helper. Define the helper once, outside the entrypoint, so repeated calls can reuse it.

Return an array of shape (x.shape[0] * int(n_repeats),). Zero repeats or an empty input give an empty array.

The diagram uses counts 2, 3, then 2 to show the same cache-reuse idea.

Loading visualization…