[Documentation] Document how to @beartype JAX dataclasses with Equinox

This works without type checking, but fails when I add the @typechecker decorator

import jax.numpy as jnp
from jax import vmap
from typing import NamedTuple, Optional, Union
from beartype import beartype as typechecker
#from typeguard import typechecked as typechecker


@typechecker
class MyParams(NamedTuple):
    mean: float

theta = MyParams(42.0)

def f(params, y):
    return params.mean + y

data = jnp.ones((10,2))
foo = vmap(f, (None,0))(theta, data)
print(foo)

It raises the error

@beartyped namedtuple_MyParams.MyParams.__new__()
 parameter mean=<object object at 0x7fb770c3a840> violates type hint <class 'float'>, 
as object <object object at 0x7fb770c3a840> not instance of float.

If I use typeguard i get the error

type of argument "mean" must be either float or int; got object instead

I think the theta variable is getting reshaped by jax to have a leading batch dimension, and this violates the type signature.

@slinderman @patrick-kidger FYI

[murphyk]

Yes, equinox does indeed solve this problem! 😃 Here is a demonstration of this, for the record (since the benefits of equinox may not be obvious to others 😃

@jaxtyped
@typechecker
#@dataclass(frozen=True) # bad
class MyParams(eqx.Module): # good
    mean: Scalar
    mean2: Float[Array, "ndim"]

theta = MyParams(1.0, jnp.ones(2))

def f(params, y):
    p = MyParams(params.mean + y, y*jnp.ones(3))
    return p

data = jnp.arange(5)

foo = vmap(lambda y: f(theta, y))(data)
print(foo)
print('mean ', foo.mean)
print('mean2 ', foo.mean2)

This type checks and produces the desired output:

MyParams(mean=f32[5], mean2=f32[5,3])
mean  [1. 2. 3. 4. 5.]
mean2
  [[0. 0. 0.]
 [1. 1. 1.]
 [2. 2. 2.]
 [3. 3. 3.]
 [4. 4. 4.]]

[patrick-kidger]

That’s because, as the error suggests, dataclasses aren’t JAX types. You should register it as PyTree. (Or use equinox.Module, which is basically @dataclass + pytree registration + a few edge-case improvements all wrapped together.)

FWIW this definitely has nothing to do with beartype, so to spare @leycec I suggest we discuss this elsewhere? 😃