Why is this function recompiling? about equinox

[lockwo]

My understanding of filter_jit is not 100%, but a common intuition I have for some problems is basically that filter_jit is like filtering by is_array, then jax.jit-ing a function with f(arr, static) and marking static as static_argnames in jax jit. What this means is that jax looks at the hashes of those functions to cache each static input. So what you have here is in the second example, the hash of the second model.forward is different (because its a totally different class initialization, I assume jax assigns it a different hash in the cache) so it triggers a recompile because jax thinks it's a different static value (when we know it isn't). If you remove the jax.jit from the fn and jit it outside (ensuring that it is the exact same cached function being used) it works without error.

import jax
import jax.numpy as jnp
import equinox as eqx
import optax
from typing import Callable
from jaxtyping import Array


class ForwardKinematicTrajectory(eqx.Module):
    trajectory: eqx.Module
    scale_mixer: Array
    forward: Callable

    def __init__(self, fk_fn: Callable):
        self.trajectory = eqx.nn.Linear(1, 3, key=jax.random.PRNGKey(0))
        self.scale_mixer = jnp.ones((3, 1))
        self.forward = fn

    def __call__(self, data):
        """Meant to be vmapped"""
        return self.forward(self.trajectory(data) * self.scale_mixer)

def loss(model: ForwardKinematicTrajectory, X, y):
    return jnp.mean(jnp.square(jax.vmap(model)(X) - y))


@eqx.filter_jit
@eqx.debug.assert_max_traces(max_traces=1)
def step(model, opt_state, X, y, loss_grad, opt):
    val, grads = loss_grad(model, X, y)
    updates, opt_state = opt.update(grads, opt_state, model)
    model = eqx.apply_updates(model, updates)
    return val, model, opt_state


# generate data
X = jnp.linspace(-1, 1, 50).reshape(-1, 1)
y = jnp.zeros((50, 3, 3))

# init model
fn = jax.jit(jnp.sin)
model = ForwardKinematicTrajectory(fn)

# optimizer
optimizer = optax.adam(1e-3)
opt_state = optimizer.init(eqx.filter(model, eqx.is_array))

# loss
loss_grad = eqx.filter_value_and_grad(loss)
loss_grad = eqx.filter_jit(loss_grad)

# this will run
for i in range(100):
    val, model, state = step(model, opt_state, X, y, loss_grad, optimizer)
    if i % 10 == 0:
        print(val)

# reinitalize model
print("reinit")
model = ForwardKinematicTrajectory(fn)
for i in range(100):
    # this will recompile the step function
    val, model, state = step(model, opt_state, X, y, loss_grad, optimizer)
    if i % 10 == 0:
        print(val)

[jdpeiffer6]

Thanks for the explanation! Makes sense and fixed my problem!