Migrate to newer nnx.Pytree api, away from flax.struct.dataclass.

Author: mathisgerdes

docs/source/intro/cg.ipynb

@@ -100,10 +100,20 @@
      "output_type": "stream",
      "text": [
       "Euler Tableau (CG1):\n",
-      " ButcherTableau(stages=1, a=((0,),), b=(1,), c=(0,))\n",
+      " \u001b[38;2;79;201;177mButcherTableau\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
+      "  \u001b[38;2;156;220;254mstages\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;182;207;169m1\u001b[0m,\n",
+      "  \u001b[38;2;156;220;254ma\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;182;207;169m0.0\u001b[0m,\u001b[38;2;255;213;3m)\u001b[0m,\u001b[38;2;255;213;3m)\u001b[0m,\n",
+      "  \u001b[38;2;156;220;254mb\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;182;207;169m1.0\u001b[0m,\u001b[38;2;255;213;3m)\u001b[0m,\n",
+      "  \u001b[38;2;156;220;254mc\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;182;207;169m0.0\u001b[0m,\u001b[38;2;255;213;3m)\u001b[0m\n",
+      "\u001b[38;2;255;213;3m)\u001b[0m\n",
       "\n",
       "CG2 Tableau:\n",
-      " ButcherTableau(stages=2, a=((0, 0), (0.5, 0)), b=(0, 1), c=(0, 0.5))\n"
+      " \u001b[38;2;79;201;177mButcherTableau\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
+      "  \u001b[38;2;156;220;254mstages\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;182;207;169m2\u001b[0m,\n",
+      "  \u001b[38;2;156;220;254ma\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;182;207;169m0.0\u001b[0m, \u001b[38;2;182;207;169m0.0\u001b[0m\u001b[38;2;255;213;3m)\u001b[0m, \u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;182;207;169m0.5\u001b[0m, \u001b[38;2;182;207;169m0.0\u001b[0m\u001b[38;2;255;213;3m)\u001b[0m\u001b[38;2;255;213;3m)\u001b[0m,\n",
+      "  \u001b[38;2;156;220;254mb\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;182;207;169m0.0\u001b[0m, \u001b[38;2;182;207;169m1.0\u001b[0m\u001b[38;2;255;213;3m)\u001b[0m,\n",
+      "  \u001b[38;2;156;220;254mc\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;182;207;169m0.0\u001b[0m, \u001b[38;2;182;207;169m0.5\u001b[0m\u001b[38;2;255;213;3m)\u001b[0m\n",
+      "\u001b[38;2;255;213;3m)\u001b[0m\n"
      ]
     }
    ],

docs/source/tutorials/scalar-theory.ipynb

@@ -22,15 +22,14 @@
 
 import functools
 import inspect
+from dataclasses import dataclass
 
-import flax
 import jax
 import jax.numpy as jnp
 import numpy as np
 from flax import nnx
 from jax_autovmap import autovmap
 
-from ..utils import Const
 from .base import Bijection
 
 
@@ -73,11 +72,6 @@ class BinaryMask(Bijection):
         >>> # reconstructed == x
     """
 
-    masks: Const
-    primary_indices: Const
-    secondary_indices: Const
-    event_shape: tuple[int, ...]
-
     def __init__(
         self,
         primary_indices: tuple[np.ndarray, ...],
@@ -90,20 +84,20 @@ def __init__(
             masks = (mask, ~mask)
         if secondary_indices is None:
             secondary_indices = np.where(masks[1])
-        self.masks = Const(masks)
-        self.primary_indices = Const(primary_indices)
-        self.secondary_indices = Const(secondary_indices)
-        self.event_shape = event_shape
+        self.primary_indices = nnx.data(primary_indices)
+        self.event_shape = nnx.static(event_shape)
+        self.masks = nnx.data(masks)
+        self.secondary_indices = nnx.data(secondary_indices)
 
     @property
     def count_primary(self):
         """Number of elements in the primary (True) mask region."""
-        return self.primary_indices.value[0].size
+        return self.primary_indices[0].size
 
     @property
     def count_secondary(self):
         """Number of elements in the secondary (False) mask region."""
-        return self.secondary_indices.value[0].size
+        return self.secondary_indices[0].size
 
     @property
     def counts(self):
@@ -145,7 +139,7 @@ def from_boolean_mask(cls, mask: jax.Array):
     @property
     def boolean_mask(self):
         """Primary boolean mask array."""
-        return self.masks.value[0]
+        return self.masks[0]
 
     def indices(
         self, extra_channel_dims: int = 0, batch_safe: bool = True, primary: bool = True
@@ -161,7 +155,7 @@ def indices(
             Indexing tuple suitable for array subscripting.
         """
         ind = (...,) if batch_safe else ()
-        ind += self.primary_indices.value if primary else self.secondary_indices.value
+        ind += self.primary_indices if primary else self.secondary_indices
         ind += (np.s_[:],) * extra_channel_dims
         return ind
 
@@ -172,10 +166,10 @@ def flip(self):
             New BinaryMask with primary and secondary regions swapped.
         """
         return self.__class__(
-            self.secondary_indices.value,
+            self.secondary_indices,
             self.event_shape,
-            masks=self.masks.value[::-1],
-            secondary_indices=self.primary_indices.value,
+            masks=self.masks[::-1],
+            secondary_indices=self.primary_indices,
         )
 
     def split(self, array, extra_channel_dims: int = 0, batch_safe: bool = True):
@@ -323,7 +317,7 @@ def checker_mask(shape, parity: bool):
     return BinaryMask.from_boolean_mask(mask.astype(bool))
 
 
-class ModuleReconstructor:
+class ModuleReconstructor(nnx.Pytree):
     """
     Parameter management utility for dynamically parameterizing modules.
 
@@ -341,42 +335,39 @@ class ModuleReconstructor:
         - Full nnx state, use `from_params`
     """
 
-    # params_treedef: Any  # static
-    # params_leaves: list[jax.core.ShapedArray]  # static
-    # unconditional: nnx.State  # array leaf
-    # graph: Any | None = None  # static
-
     def __init__(
-        self, module_or_state: nnx.State | nnx.Module, filter: nnx.Param = nnx.Param
+        self,
+        module_or_state: nnx.State | nnx.Module,
+        filter: nnx.Param = nnx.Param,
     ):
         if isinstance(module_or_state, nnx.State):
             self.graph = None
             state = module_or_state
         else:
             graph, state = nnx.split(module_or_state)
-            self.graph = graph
+            self.graph = nnx.static(graph)
 
         params, unconditional = nnx.split_state(state, filter, ...)
 
         params = jax.tree.map(lambda x: jax.core.ShapedArray(x.shape, x.dtype), params)
 
         params_leaves, params_treedef = jax.tree.flatten(params)
 
-        self.params_treedef = params_treedef
-        self.params_leaves = params_leaves
-        self.unconditional = unconditional
+        self.params_treedef = nnx.static(params_treedef)
+        self.params_leaves = nnx.static(params_leaves)
+        self.unconditional = nnx.data(unconditional)
 
-    def _tree_flatten(self):
-        children = (self.unconditional,)
-        aux_data = (self.params_treedef, self.params_leaves, self.graph)
-        return children, aux_data
+    # def _tree_flatten(self):
+    #     children = (self.unconditional,)
+    #     aux_data = (self.params_treedef, self.params_leaves, self.graph)
+    #     return children, aux_data
 
-    @classmethod
-    def _tree_unflatten(cls, aux_data, children):
-        self = object.__new__(cls)
-        self.params_treedef, self.params_leaves, self.graph = aux_data
-        (self.unconditional,) = children
-        return self
+    # @classmethod
+    # def _tree_unflatten(cls, aux_data, children):
+    #     self = object.__new__(cls)
+    #     self.params_treedef, self.params_leaves, self.graph = aux_data
+    #     (self.unconditional,) = children
+    #     return self
 
     @property
     def params(self):
@@ -490,15 +481,15 @@ def __repr__(self):
         return f"ModuleReconstructor:{state_or_module}"
 
 
-jax.tree_util.register_pytree_node(
-    ModuleReconstructor,
-    ModuleReconstructor._tree_flatten,
-    ModuleReconstructor._tree_unflatten,
-)
+# jax.tree_util.register_pytree_node(
+#     ModuleReconstructor,
+#     ModuleReconstructor._tree_flatten,
+#     ModuleReconstructor._tree_unflatten,
+# )
 
 
-@flax.struct.dataclass
-class AutoVmapReconstructor:
+@dataclass(frozen=True)
+class AutoVmapReconstructor(nnx.Pytree):
     r"""Automatic vectorization for module reconstruction with batched parameters.
 
     This class provides a solution for bijections that do not natively support
@@ -539,8 +530,8 @@ class AutoVmapReconstructor:
     """
 
     reconstructor: ModuleReconstructor
-    params: nnx.State | dict | list[jax.Array] | jax.Array
-    params_rank: int | dict = 1
+    params: nnx.Data[nnx.State | dict | list[jax.Array] | jax.Array]
+    params_rank: nnx.Data[int | dict] = 1
 
     def __call__(self, fn_name, *args, input_ranks: tuple[int, ...] = (0, 0), **kwargs):
 

src/bijx/bijections/coupling.py

@@ -20,15 +20,14 @@
 
 from functools import partial, reduce
 
-import flax.struct
 import jax
 import jax.numpy as jnp
 import numpy as np
+from flax import nnx
 from jax import core, custom_derivatives
 
 
-@flax.struct.dataclass
-class ButcherTableau:
+class ButcherTableau(nnx.Pytree):
     r"""Butcher tableau defining coefficients for Runge-Kutta integration schemes.
 
     Encodes the coefficient structure for explicit Runge-Kutta methods in
@@ -50,29 +49,16 @@ class ButcherTableau:
     Consistency requires: $c_i = \sum_j a_{ij}$ and $\sum_i b_i = 1$.
     """
 
-    stages: int
-    """Number of stages $s$ in the method."""
-
-    a: tuple[tuple[int, ...]]
-    """Coefficient matrix $(a_{ij})$ as nested tuples."""
-
-    b: tuple[int, ...]
-    """Weight vector $(b_i)$ as tuple."""
-
-    c: tuple[int, ...]
-    """Node vector $(c_i)$ as tuple (computed from $a$)."""
-
-    @classmethod
-    def from_ab(cls, a, b):
+    def __init__(self, a, b):
         r"""Construct Butcher tableau from coefficient matrix and weights.
 
         Creates a ButcherTableau instance from the $a$ matrix and $b$ vector,
         automatically computing the node vector $c_i = \sum_j a_{ij}$ and
         validating consistency conditions.
 
         Args:
-            a: Coefficient matrix as list of lists, shape $(s, s)$.
-            b: Weight vector as list, length $s$.
+            a: Coefficient matrix or list of lists, shape $(s, s)$.
+            b: Weight vector or list, length $s$.
 
         Returns:
             ButcherTableau instance with computed node vector.
@@ -85,12 +71,12 @@ def from_ab(cls, a, b):
 
         Example:
             >>> # Second-order Crouch-Grossmann method
-            >>> cg2 = ButcherTableau.from_ab(
+            >>> cg2 = ButcherTableau(
             ...     a=[[0, 0], [1/2, 0]], b=[0, 1]
             ... )
         """
-        a = tuple(tuple(ai) for ai in a)
-        b = tuple(b)
+        a = tuple(tuple(float(aij) for aij in ai) for ai in a)
+        b = tuple(float(bi) for bi in b)
         c = tuple(sum(ai) for ai in a)
 
         assert all(len(ai) == len(c) for ai in a)
@@ -101,10 +87,13 @@ def from_ab(cls, a, b):
             for i in range(j + 1):
                 assert a[i][j] == 0, "only explicit methods supported"
 
-        return cls(stages=len(c), a=a, b=b, c=c)
+        self.stages = len(c)
+        self.a = nnx.static(a)
+        self.b = nnx.static(b)
+        self.c = nnx.static(c)
 
 
-EULER = ButcherTableau.from_ab(
+EULER = ButcherTableau(
     a=[[0]],
     b=[1],
 )
@@ -115,7 +104,7 @@ def from_ab(cls, a, b):
 For Lie groups: $g_{n+1} = \exp(h A(t_n, g_n)) g_n$.
 """
 
-CG2 = ButcherTableau.from_ab(
+CG2 = ButcherTableau(
     a=[[0, 0], [1 / 2, 0]],
     b=[0, 1],
 )
@@ -131,7 +120,7 @@ def from_ab(cls, a, b):
 Update: $g_{n+1} = \exp(h k_2) g_n$
 """
 
-CG3 = ButcherTableau.from_ab(
+CG3 = ButcherTableau(
     a=[[0, 0, 0], [3 / 4, 0, 0], [119 / 216, 17 / 108, 0]],
     b=[13 / 51, -2 / 3, 24 / 17],
 )
@@ -432,3 +421,4 @@ def augmented_ode(t, state, args):
 
 
 _crouch_grossmann.defvjp(_crouch_grossmann_fwd, _crouch_grossmann_rev)
+_crouch_grossmann.defvjp(_crouch_grossmann_fwd, _crouch_grossmann_rev)