Force a dynamic variable to be static

[chrisflesher]

Is there a way to tell the JAX compiler, "I know it's not supposed to be static but just make it static anyways!"

Below is a simplified example of the issue:

import typing
import jax
import jax.numpy as jnp

class Image(typing.NamedTuple):
    resolution: float
    values: jax.Array

@jax.jit
def sample_image(image: Image) -> jax.Array:
    num_samples = int(1/image.resolution)  # <- want to force this to be static
    return image.values * num_samples

num_images = 3
num_values_per_image = 5
static_resolution = 0.2
images = jax.vmap(Image)(jnp.full(num_images, static_resolution), jnp.ones((num_images, num_values_per_image)))
sampled_values = jax.vmap(sample_image)(images)

Running this causes a jax.errors.ConcretizationTypeError because image.resolution is not static. To fix this I have to add an extra static argument:

@functools.partial(jax.jit, static_argnames=['resolution'])
def sample_values(image: Image, resolution: float) -> jax.Array:
    num_samples = int(1/resolution)  # works
    return image.values * num_samples

However it can be cumbersome to add extra arguments like this functions buried several layers deep when I know resolution won't ever change. Is there a way to force the JAX compiler to use the first instance of dynamic value as static?

[jakevdp]

The way to force a variable to be static within a jit-compiled function is using static_argnums, as mentioned in the jax.jit docs.

Unfortunately, in your particular case, because you are wrapping the function in vmap, there is no way to mark the input as static: a vmapped value is always dynamic, by definition.

To fix this you'll need to do two things:

1. adjust your vmap call so you are not mapping over the value you want to be static
2. adjust your jit call so that you mark the particular value as static via static_argnames. One way to do this is by jit-compiling a sub-computation in which you pass the values directly to the function.

Doing these together might look like this:

from functools import partial
import typing
import jax
import jax.numpy as jnp

class Image(typing.NamedTuple):
    resolution: jax.Array
    values: jax.Array

def sample_image(image: Image) -> jax.Array:
    return _sample_image(image.resolution, image.values)

@partial(jax.jit, static_argnames=["resolution"])
def _sample_image(resolution, values):
    num_samples = int(1/resolution)
    return values * num_samples

num_images = 3
images = Image(0.2, jnp.ones((num_images, 5)))
sampled_values = jax.vmap(sample_image, in_axes=(Image(None, 0),))(images)

If you are using Image frequently within jit-compiled functions and you want to keep resolution static in a more automatic fashion, you could register your Image object as a custom pytree, where you can specify static components by including them in aux_data (see Extending Pytrees):

from jax.tree_util import register_pytree_node

register_pytree_node(
    Image,
    lambda image: ([image.values], [image.resolution]),
    lambda aux_data, children: Image(resolution=aux_data[0], values=children[0])
)

@jax.jit
def sample_image(image: Image) -> jax.Array:
    num_samples = int(1/image.resolution)
    return image.values * num_samples

sample_image(Image(0.2, jnp.arange(5)))  # now resolution is always treated as static

[chrisflesher]

Think aux_data is exactly what I'm looking for, thanks!