Add flatten mode to improve tree_flatten and tree_unflatten speed when needed.

Author: NTT123

Diff for: opax/transform.py

@@ -18,7 +18,7 @@ def __init__(self, params=None):
     def __call__(self, updates, params=None):
         raise NotImplementedError("A subclass must implement this method")
 
-    def step(self, grads, params, all_finite: Optional[bool] = None):
+    def step(self, grads, params, all_finite: Optional[jnp.ndarray] = None):
         """An optimizing step.
 
         First, transform gradients
@@ -291,15 +291,36 @@ def __call__(self, updates, params=None):
 def chain(*fs: Callable[[Any], GradientTransformation]):
     class Chain(GradientTransformation):
         transforms: Sequence[GradientTransformation]
+        flatten: bool
 
-        def __init__(self, params):
+        def __init__(self, params, flatten: bool = False):
+            """Create a chain of gradient transformations.
+
+            Arguments:
+                params: trainable parameters.
+                flatten: flatten trainable parameters to a list for faster speed in jit mode.
+            """
             super().__init__()
-            transforms = [f(params) for f in fs]
+            self.flatten = flatten
+            if flatten:
+                leaves = jax.tree_leaves(params)
+                transforms = [f(leaves) for f in fs]
+            else:
+                transforms = [f(params) for f in fs]
             self.register_module_subtree("transforms", transforms)
 
         def __call__(self, updates, params=None):
-            for f in self.transforms:
-                updates = f(updates=updates, params=params)
+            if self.flatten:
+                updates_leaves, updates_treedef = jax.tree_flatten(updates)
+                params_leaves = jax.tree_leaves(params)
+
+                for f in self.transforms:
+                    updates_leaves = f(updates=updates_leaves, params=params_leaves)
+
+                updates = jax.tree_unflatten(updates_treedef, updates_leaves)
+            else:
+                for f in self.transforms:
+                    updates = f(updates=updates, params=params)
 
             return updates
 

Diff for: setup.py

@@ -1,6 +1,6 @@
 from setuptools import find_packages, setup
 
-__version__ = "0.1.6"
+__version__ = "0.1.7"
 url = "https://github.com/ntt123/opax"
 
 install_requires = ["pax"]

Diff for: tests/test_opax.py

@@ -112,3 +112,20 @@ def loss_fn(params, model, inputs) -> pax.utils.LossFnOutput:
     with pytest.raises(ValueError):
         for i in range(10):
             loss, net, opt = update_fn(net, opt, (x, x))
+
+
+def test_train_flatten():
+    net = pax.nn.Sequential(
+        pax.nn.Linear(1, 2),
+        pax.nn.Linear(2, 1),
+    )
+
+    def loss_fn(params, model, inputs) -> pax.utils.LossFnOutput:
+        loss = jnp.mean(jnp.square(model.update(params)(inputs[0]) - inputs[1]))
+        return loss, (loss, model)
+
+    update_fn = pax.utils.build_update_fn(loss_fn)
+    x = jnp.zeros((1, 1))
+    opt = opax.adam()(net.parameters(), flatten=True)
+    for i in range(10):
+        loss, net, opt = update_fn(net, opt, (x, x))