JAX + dataclasses

I have a rather big dataclass to describe a robot model, that includes a graph of links and a list of joints. Each node of the graph references the parent link and all the child links. Each joint object references its parent and child links.

When I try to copy_and_mutate any of these objects, maybe due to all this nesting, an infinite loop occurs. I suspect that the existing logic tries to unfreeze all the leafs of the pytree, but the high interconnection and the properties of mutable Python types lead to a never ending unfreezing process.

This PR addresses this edge case by storing the list of IDs of objects already unfreezed. It solves my problem, and it should not add any noticeable performance degradation.

cc @brentyi

First of all, thank you for the amazing library! I have recently discovered jax_dataclasses and I have decided to port my messy JAX functional code to a more organised object-oriented code based on jax_dataclasses.

In my application, I have some derived quantities of the attributes of the dataclass that are static values used to determine the shape of tensors during JIT compilation. I would like to include them as attribute of the dataclass, but I'm getting an error and I would like to know if there is workaround.

Here is a simple example, where the attribute _sum is a derived static field that depends on the constant value of the array a.

import jax
import jax.numpy as jnp
import jax_dataclasses as jdc

@jdc.pytree_dataclass()
class PyTreeDataclass:
    a: jnp.ndarray
    _sum: int = jdc.static_field(init=False, repr=False)

    def __post_init__(self):
        object.__setattr__(self, "_sum", self.a.sum().item())

def print_pytree(obj):
    print(obj._sum)

obj = PyTreeDataclass(jnp.arange(4))
print_pytree(obj)
jax.jit(print_pytree)(obj)

The non-jitted version works, but when print_pytree is jitted I get the following error.

File "jax_dataclasses_issue.py", line 14, in __post_init__
    object.__setattr__(self, "_sum", self.a.sum().item())
AttributeError: 'bool' object has no attribute 'sum'

Is there a way to compute in the __post_init__ the value of static fields not initialized in __init__ that depend on jnp.ndarray attributes of the dataclass?

The file jax_dataclasses/_copy_and_mutate.py raises many warnings complaining a deprecated function.

FutureWarning: jax.tree_leaves is deprecated, and will be removed in a future release. Use jax.tree_util.tree_leaves instead.

I just discovered the jaxtyping library and I think it could be an interesting alternative to the current typing system proposed by jax_dataclasses.

jaxtyping supports variable-size axes and symbolic expressions in terms of other variable-size axes, see https://github.com/google/jaxtyping/blob/main/API.md and it has very few requirements.

Do you think that it could be added to jax_dataclasses?

Thanks for the handy library!

I have a pytree_dataclass that contains a few static_fields that I would like to have serialized by the facilities in flax.serialize. I noticed that jax_dataclasses.asdict handles these, but that flax.serialization.to_state_dict and flax.serialization.to_bytes both ignore them. What is the correct way (if any) to have these fields included in flax's serialization? Should I be using another technique?

import jax_dataclasses as jdc
from jax import numpy as jnp
import flax.serialization as fs


@jdc.pytree_dataclass
class Demo:
    a: jnp.ndarray = jnp.ones(3)
    b: bool = jdc.static_field(default=False)


demo = Demo()
print(f'{jdc.asdict(demo) = }')
print(f'{fs.to_state_dict(demo) = }')
print(f'{fs.from_bytes(Demo, fs.to_bytes(demo)) = }')

# jdc.asdict(demo) = {'a': array([1., 1., 1.]), 'b': False}
# fs.to_state_dict(demo) = {'a': DeviceArray([1., 1., 1.], dtype=float64)}
# fs.from_bytes(Demo, fs.to_bytes(demo)) = {'a': array([1., 1., 1.])}

Thanks in advance!