Why does jitting fail when passing model as argument?

[Rasmuskh]

Hi,

I noticed that jit compilation fails when my model is passed as a parameter to a function, but not when it is just used from within the outer scope.

Here is a minimal working example:

import jax, flax

@jax.jit
def infer_v1(params, x):
    y = model.apply({'params': params}, x, method=model.__call__)
    return y

@jax.jit
def infer_v2(params, m, x):
    y = m.apply({'params': params}, x, method=m.__call__)
    return y

key1, key2 = jax.random.PRNGKey(102), jax.random.PRNGKey(774)
x = jax.random.normal(key1, shape=([1, 3]), dtype=jax.numpy.float32)
model = flax.linen.Dense(2)
params = model.init(key2, x)['params']

y1 = infer_v1(params, x)
y2 = infer_v2(params, model, x)

jit compiling works fine for infer_v1, but leads to an error for infer_v2. The error ends with "AttributeError: DynamicJaxprTracer has no attribute apply".

If I remove @jax.jit from infer_v2, then it works as expected. What is going on here? Why does jitting fail when the model is passed as an argument?

(Additional context: I normally create a flax TrainState and use state.apply_fn instead of directly calling model.apply, but I am dealing with a case where I need to define and call multiple class methods for my network during training, which state.apply_fn doesn't let me do.)

[chiamp]

jax.jit won't work on a function with a Module argument type. We can demonstrate this by making a dummy function with the same argument signature:

@jax.jit
def infer_v2(params, m, x):
    return

key1, key2 = jax.random.PRNGKey(102), jax.random.PRNGKey(774)
x = jax.random.normal(key1, shape=([1, 3]), dtype=jax.numpy.float32)
model = flax.linen.Dense(2)
params = model.init(key2, x)['params']

y2 = infer_v2(params, model, x)

We get an error:

<class 'flax.linen.linear.Dense'> is not a valid JAX type.

One solution to this is you can specify that that argument is static (check out the API documentation on static_argnums here) via:

def infer_v2(params, m, x):
    y = m.apply({'params': params}, x, method=m.__call__)
    return y

key1, key2 = jax.random.PRNGKey(102), jax.random.PRNGKey(774)
x = jax.random.normal(key1, shape=([1, 3]), dtype=jax.numpy.float32)
model = flax.linen.Dense(2)
params = model.init(key2, x)['params']

y2 = jax.jit(infer_v2, static_argnums=1)(params, model, x)

But jax.jit will trigger a recompile every time you pass in a different value for m (read more about it here).

Alternatively, you could pass in a string to the method argument in .apply to call the specific method you want:

import jax, jax.numpy as jnp
import flax, flax.linen as nn
from flax.training import train_state
import optax

class MultiMethodModule(nn.Module):
  @nn.compact
  def __call__(self, x):
    return nn.Dense(2)(x)
  def other_call(self, x):
    return self(x) + 1

x = jnp.ones((1, 3))
model = MultiMethodModule()
variables = model.init(jax.random.key(0), x)
state = train_state.TrainState.create(apply_fn=model.apply, params=variables['params'], tx=optax.adam(1e-3))

@jax.jit
def infer_multi_method(state, x):
  call_out = state.apply_fn({'params': state.params}, x, method='__call__')
  other_call_out = state.apply_fn({'params': state.params}, x, method='other_call')
  return call_out, other_call_out

call_out, other_call_out = infer_multi_method(state, x)
print(call_out)
print(other_call_out)

[[-1.4500693  1.3251126]]
[[-0.4500693  2.3251126]]

Flax also provides "lifted transformed" versions of JAX transformations, which would allow you to use Module arguments. You can view the full list here.

But I think for your situation, one of the above solutions should suffice.

[Rasmuskh]

Thanks for the thorough answer!

Good to know that jit doesn't work with nn.module objects as arguments (unless set as static arguments).
For my use case the best solution was your suggestion of providing the method keyword argument to state.apply_fn.

Thanks for the help! 😃