Weird error when JITting a method

[IrishWhiskey]

Consider the following code:

import jax
import jax.numpy as jnp

@jax.tree_util.register_pytree_node_class
class A:
    def __init__(self, eps: int):
        self.eps = eps
    
    @jax.jit
    def loss(self, X, w):
        return jnp.dot(X, w) * self.eps
    
    def tree_flatten(self):
        children = ()  # arrays / dynamic values
        aux_data = dict(eps=self.eps)  # static values
        return children, aux_data

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        return cls(*children, **aux_data)

X = jax.random.normal(jax.random.PRNGKey(0), shape=(1000, 5))
w = jax.random.normal(jax.random.PRNGKey(0), shape=(5,))

a = A(jnp.sum(w))
jax.make_jaxpr(a.loss)(X, w)

a = A(jnp.sum(w))
jax.make_jaxpr(a.loss)(X, w)

It looks pretty innocuous but when I run it the last line raises an error. The error message is the following:

jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/ec2-user/SageMaker/main.py", line 29, in <module>
    jax.make_jaxpr(a.loss)(X, w)
jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[]..
The error occurred while tracing the function loss at /home/ec2-user/SageMaker/main.py:9 for make_jaxpr. This value became a tracer due to JAX operations on these lines:

  operation a:bool[] = eq b c
    from line /home/ec2-user/SageMaker/main.py:29 (<module>)
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError

I noticed that if I remove the jit of the loss method the error disappears. Why is that? Am I using it wrongly?

I am using jax 0.4.23 and python 3.10 on a SageMaker ml.p3 instance.

[jakevdp]

We have a section in the FAQ about using jit with class methods: https://jax.readthedocs.io/en/latest/faq.html#how-to-use-jit-with-methods

Please let me know if that doesn't answer your question!

[jakevdp]

Looking closer, I think this is a problem with pytree registration. In your code eps is a jax array, but you're flattening as if it's a static argument. There are two ways to fix this: first, you could keep it as a dynamic argument and include eps in the children rather than the aux_data:

    def tree_flatten(self):
        children = (self.eps,)  # arrays / dynamic values
        aux_data = {}  # static values
        return children, aux_data

Alternatively, you could store it as an int rather than an array, then it could be a true static value:

    def __init__(self, eps: int):
        self.eps = int(eps)