No error raised by jitting class methods with backend agnostic architecture

[king-p3nguin]

Hi, I am writing a code using backend-agnostic architecture so that I can run my existing code in both Numpy and Jax. The code is something like this:

import numpy as np
import jax.numpy as jnp
from jax import jit
import jax

class NumPyBackend:
    def sin(self, x):
        return np.sin(x)
    def cos(self, x):
        return np.cos(x)
    def jit(self, f):
        return f


class JaxBackend:
    def sin(self, x):
        return jnp.sin(x)
    def cos(self, x):
        return jnp.cos(x)
    def jit(self, f):
        return jit(f)

BACKEND = NumPyBackend()


class CustomClass:
    def __init__(self, x):
        self.x = x

    @BACKEND.jit
    def calc(self):
        if jnp.isclose(BACKEND.sin(self.x)**2 + BACKEND.cos(self.x)**2, 1):
            return BACKEND.sin(self.x)**2 + BACKEND.cos(self.x)**2

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

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

from jax import tree_util
tree_util.register_pytree_node(CustomClass,
                               CustomClass._tree_flatten,
                               CustomClass._tree_unflatten)

a = 1.0

BACKEND = NumPyBackend()
a_np = CustomClass(a)

print(a_np.calc())
# 1.0 (No error)

a = jnp.array(a)

BACKEND = JaxBackend()
a_jax = CustomClass(a)

print(a_jax.calc())
# 1.0 (No error)

@jit
def a_ref(x):
    if jnp.isclose(BACKEND.sin(x)**2 + BACKEND.cos(x)**2, 1):
        return jnp.sin(x)**2 + jnp.cos(x)**2

print(a_ref(a))
# TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[]..
# The error occurred while tracing the function a_ref at /tmp/ipykernel_14657/129821386.py:62 for jit. This concrete value was not available in Python because it depends on the value of the argument x.
# See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError

I am wondering why print(a_jax.calc()) did not raise TracerBoolConversionError whereas print(a_ref(a)) did. I would like to know the reason why this happens and how to improve the code so that the error appears correctly. I appreciate your help.

Note that the following code will raise TracerBoolConversionError.

import jax

class CustomClass:
    def __init__(self, x):
        self.x = x

    @jax.jit
    def norm(self):
        if jnp.isclose(jnp.sin(self.x)**2 + jnp.cos(self.x)**2, 1):
            return jnp.sin(self.x)**2 + jnp.cos(self.x)**2

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

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

from jax import tree_util
tree_util.register_pytree_node(CustomClass,
                               CustomClass._tree_flatten,
                               CustomClass._tree_unflatten)

a_ref2 = CustomClass(a)
a_ref2.norm()
# TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[]..
# The error occurred while tracing the function norm at /tmp/ipykernel_20399/589877195.py:7 for jit. This concrete value was not available in Python because it depends on the value of the argument self[<flat index 0>].
# See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError

OS: Windows WSL2 (Ubuntu)
Python 3.10.12
jax==0.4.23
jaxlib==0.4.23
numpy==1.25.2

[jakevdp]

Hi - the reason that a_jax.calc() does not error is because the calc method of a_jax is not JIT-compiled. Class method decorators are evaluated at the time the class is defined, not at the time the class is instantiated, and when the class is defined, the global variable BACKEND is set to NumPyBackend(). The fact that this global variable is modified later has no effect on the @BACKEND.jit() decorator used previously in the class definition.

[king-p3nguin]

Thank you for the answer! I tried BACKEND.jit(a_jax.calc) and this seems to be working now.

[patrick-kidger]

Side note -- note that registering classes as custom PyTrees can be fairly finickity business in JAX. You might like Equinox, which offers an equinox.Module designed to make PyTree-registration-of-classes easier.