brainstate 0.0.2.post20241009__py2.py3-none-any.whl → 0.1.0__py2.py3-none-any.whl

brainstate/optim/_lr_scheduler_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
+from __future__ import annotations
 
 import unittest
 
@@ -22,28 +23,28 @@ import brainstate as bst
 
 
 class TestMultiStepLR(unittest.TestCase):
-  def test1(self):
-    lr = bst.optim.MultiStepLR(0.1, [10, 20, 30], gamma=0.1)
-    for i in range(40):
-      r = lr(i)
-      if i < 10:
-        self.assertEqual(r, 0.1)
-      elif i < 20:
-        self.assertTrue(jnp.allclose(r, 0.01))
-      elif i < 30:
-        self.assertTrue(jnp.allclose(r, 0.001))
-      else:
-        self.assertTrue(jnp.allclose(r, 0.0001))
+    def test1(self):
+        lr = bst.optim.MultiStepLR(0.1, [10, 20, 30], gamma=0.1)
+        for i in range(40):
+            r = lr(i)
+            if i < 10:
+                self.assertEqual(r, 0.1)
+            elif i < 20:
+                self.assertTrue(jnp.allclose(r, 0.01))
+            elif i < 30:
+                self.assertTrue(jnp.allclose(r, 0.001))
+            else:
+                self.assertTrue(jnp.allclose(r, 0.0001))
 
-  def test2(self):
-    lr = bst.transform.jit(bst.optim.MultiStepLR(0.1, [10, 20, 30], gamma=0.1))
-    for i in range(40):
-      r = lr(i)
-      if i < 10:
-        self.assertEqual(r, 0.1)
-      elif i < 20:
-        self.assertTrue(jnp.allclose(r, 0.01))
-      elif i < 30:
-        self.assertTrue(jnp.allclose(r, 0.001))
-      else:
-        self.assertTrue(jnp.allclose(r, 0.0001))
+    def test2(self):
+        lr = bst.transform.jit(bst.optim.MultiStepLR(0.1, [10, 20, 30], gamma=0.1))
+        for i in range(40):
+            r = lr(i)
+            if i < 10:
+                self.assertEqual(r, 0.1)
+            elif i < 20:
+                self.assertTrue(jnp.allclose(r, 0.01))
+            elif i < 30:
+                self.assertTrue(jnp.allclose(r, 0.001))
+            else:
+                self.assertTrue(jnp.allclose(r, 0.0001))

brainstate/optim/_optax_optimizer.py

@@ -17,192 +17,119 @@
 from __future__ import annotations
 
 import importlib.util
-from typing import Any
+from typing import Hashable, Dict, Optional
 
-import jax.numpy as jnp
+from brainstate._state import ShortTermState, State, StateDictManager
+from brainstate.typing import PyTree
+from ._base import Optimizer
 
-from brainstate._module import Module
-from brainstate._state import ShortTermState, ParamState
+optax_installed = importlib.util.find_spec('optax') is not None
 
 __all__ = [
-  'OptaxOptimizer',
+    'OptaxOptimizer',
 ]
 
-optax_installed = importlib.util.find_spec('optax') is not None
 
+class OptaxOptimizer(Optimizer):
+    """Simple train state for the common case with a single Optax optimizer.
 
-class OptaxState(ShortTermState):
-  """Wrapper class for Optimizer Variables."""
-  pass
-
-
-class OptaxOptimizer(Module):
-  """Simple train state for the common case with a single Optax optimizer.
-
-  Example usage::
-
-    >>> import jax, jax.numpy as jnp
-    >>> import brainstate as bst
-    >>> from brainstate import nn
-    >>> import optax
-    ...
-    >>> class Model(bst.Module):
-    ...   def __init__(self):
-    ...     super().__init__()
-    ...     self.linear1 = nn.Linear(2, 3)
-    ...     self.linear2 = nn.Linear(3, 4)
-    ...   def __call__(self, x):
-    ...     return self.linear2(self.linear1(x))
-    ...
-    >>> x = jax.random.normal(jax.random.key(0), (1, 2))
-    >>> y = jnp.ones((1, 4))
-    ...
-    >>> model = Model()
-    >>> tx = optax.adam(1e-3)
-    >>> state = bst.optim.OptaxOptimizer(model, tx)
-    ...
-    >>> loss_fn = lambda model: ((model(x) - y) ** 2).mean()
-    >>> loss_fn(model)
-    Array(1.7055722, dtype=float32)
-    >>> grads = bst.transform.grad(loss_fn)(state.model)
-    >>> state.update(grads)
-    >>> loss_fn(model)
-    Array(1.6925814, dtype=float32)
-
-  Note that you can easily extend this class by subclassing it for storing
-  additional data (e.g. adding metrics).
-
-  Example usage::
-
-    >>> class TrainState(nnx.Optimizer):
-    ...   def __init__(self, model, tx, metrics):
-    ...     self.metrics = metrics
-    ...     super().__init__(model, tx)
-    ...   def update(self, *, grads, **updates):
-    ...     self.metrics.update(**updates)
-    ...     super().update(grads)
-    ...
-    >>> metrics = nnx.metrics.Average()
-    >>> state = TrainState(model, tx, metrics)
-    ...
-    >>> grads = nnx.grad(loss_fn)(state.model)
-    >>> state.update(grads=grads, values=loss_fn(state.model))
-    >>> state.metrics.compute()
-    Array(1.6925814, dtype=float32)
-    >>> state.update(grads=grads, values=loss_fn(state.model))
-    >>> state.metrics.compute()
-    Array(1.68612, dtype=float32)
-
-  For more exotic usecases (e.g. multiple optimizers) it's probably best to
-  fork the class and modify it.
-
-  Attributes:
-    step: An ``OptaxState`` :class:`Variable` that tracks the step count.
-    model: The wrapped :class:`Module`.
-    tx: An Optax gradient transformation.
-    opt_state: The Optax optimizer state.
-  """
-
-  def __init__(
-      self,
-      model: Module,
-      tx: 'optax.GradientTransformation',
-      wrt: Any = ParamState,
-  ):
-    """
-    Instantiate the class and wrap the :class:`Module` and Optax gradient
-    transformation. Instantiate the optimizer state to keep track of
-    :class:`Variable` types specified in ``wrt``. Set the step count to 0.
-
-    Args:
-      model: An NNX Module.
-      tx: An Optax gradient transformation.
-      wrt: optional argument to filter for which :class:`Variable`'s to keep
-        track of in the optimizer state. These should be the :class:`Variable`'s
-        that you plan on updating; i.e. this argument value should match the
-        ``wrt``  argument passed to the ``nnx.grad`` call that will generate the
-        gradients that will be passed into the ``grads`` argument of the
-        :func:`update` method.
-    """
+    Example usage::
 
-    # tx must be an instance of optax.GradientTransformation
-    import optax  # type: ignore[import-not-found,import-untyped]
-    if not isinstance(tx, optax.GradientTransformation):
-      raise TypeError(f"tx must be an instance of optax.GradientTransformation, got {tx}")
-    self.tx = tx
-
-    # model
-    if not callable(model):
-      raise TypeError(f"model must be a callable, got {model}")
-    self.model = model
-
-    # wrt
-    self.opt_state = tx.init(nnx.state(model, wrt))
-    self.wrt = wrt
-
-  def update(self, grads):
-    """Updates ``step``, ``params``, ``opt_state`` and ``**kwargs`` in return value.
-    The ``grads`` must be derived from ``nnx.grad(..., wrt=self.wrt)``, where the
-    gradients are with respect to the same :class:`Variable` types as defined in
-    ``self.wrt`` during instantiation of this ``Optimizer``. For example::
-
-      >>> from flax import nnx
-      >>> import jax, jax.numpy as jnp
+      >>> import jax
+      >>> import jax.numpy as jnp
+      >>> import brainstate as bst
       >>> import optax
-
-      >>> class CustomVariable(nnx.Variable):
-      ...   pass
-
-      >>> class Model(nnx.Module):
-      ...   def __init__(self, rngs):
-      ...     self.linear = nnx.Linear(2, 3, rngs=rngs)
-      ...     self.custom_variable = CustomVariable(jnp.ones((1, 3)))
+      ...
+      >>> class Model(bst.nn.Module):
+      ...   def __init__(self):
+      ...     super().__init__()
+      ...     self.linear1 = bst.nn.Linear(2, 3)
+      ...     self.linear2 = bst.nn.Linear(3, 4)
       ...   def __call__(self, x):
-      ...     return self.linear(x) + self.custom_variable
-      >>> model = Model(rngs=nnx.Rngs(0))
-      >>> jax.tree.map(jnp.shape, nnx.state(model))
-      State({
-        'custom_variable': VariableState(
-          type=CustomVariable,
-          value=(1, 3)
-        ),
-        'linear': {
-          'bias': VariableState(
-            type=Param,
-            value=(3,)
-          ),
-          'kernel': VariableState(
-            type=Param,
-            value=(2, 3)
-          )
-        }
-      })
-
-      >>> # update:
-      >>> # - only Linear layer parameters
-      >>> # - only CustomVariable parameters
-      >>> # - both Linear layer and CustomVariable parameters
-      >>> loss_fn = lambda model, x, y: ((model(x) - y) ** 2).mean()
-      >>> for variable in (nnx.Param, CustomVariable, (nnx.Param, CustomVariable)):
-      ...   # make sure `wrt` arguments match for `nnx.Optimizer` and `nnx.grad`
-      ...   state = nnx.Optimizer(model, optax.adam(1e-3), wrt=variable)
-      ...   grads = nnx.grad(loss_fn, argnums=nnx.DiffState(0, variable))(
-      ...     state.model, jnp.ones((1, 2)), jnp.ones((1, 3))
-      ...   )
-      ...   state.update(grads=grads)
-
-    Note that internally this function calls ``.tx.update()`` followed by a call
-    to ``optax.apply_updates()`` to update ``params`` and ``opt_state``.
-
-    Args:
-      grads: the gradients derived from ``nnx.grad``.
+      ...     return self.linear2(self.linear1(x))
+      ...
+      >>> x = bst.random.randn(1, 2)
+      >>> y = jnp.ones((1, 4))
+      ...
+      >>> model = Model()
+      >>> tx = optax.adam(1e-3)
+      >>> optimizer = bst.optim.OptaxOptimizer(tx)
+      >>> optimizer.register_trainable_weights(model.states(bst.ParamState))
+      ...
+      >>> loss_fn = lambda: ((model(x) - y) ** 2).mean()
+      >>> loss_fn()
+      Array(1.7055722, dtype=float32)
+      >>> grads = bst.augment.grad(loss_fn, model.states(bst.ParamState))()
+      >>> optimizer.update(grads)
+      >>> loss_fn()
+      Array(1.6925814, dtype=float32)
+
+    For more exotic usecases (e.g. multiple optimizers) it's probably best to
+    fork the class and modify it.
+
+    Attributes:
+      param_states: The parameter states to update.
+      tx: An Optax gradient transformation.
     """
-    import optax  # type: ignore[import-not-found,import-untyped]
-    state = nnx.state(self.model, self.wrt)
-
-    updates, new_opt_state = self.tx.update(grads, self.opt_state, state)
-    new_params = optax.apply_updates(state, updates)
-    assert isinstance(new_params, nnx.State)
 
-    nnx.update(self.model, new_params)
-    self.opt_state = new_opt_state
+    param_states: StateDictManager
+    opt_state: Optional[ShortTermState]
+
+    def __init__(
+        self,
+        tx: 'optax.GradientTransformation',
+    ):
+        """
+        Instantiate the class and wrap the :class:`FlattedDict` and Optax gradient
+        transformation. Instantiate the optimizer state to keep track of
+        :class:`State`.
+
+        Args:
+          tx: An Optax gradient transformation.
+        """
+        super().__init__()
+
+        # tx must be an instance of optax.GradientTransformation
+        import optax  # type: ignore[import-not-found,import-untyped]
+        if not isinstance(tx, optax.GradientTransformation):
+            raise TypeError(f"tx must be an instance of optax.GradientTransformation, got {tx}")
+        self.tx = tx
+
+        # optimizer state
+        self.opt_state = None
+
+    def register_trainable_weights(self, param_states: Dict[Hashable, State]):
+        # model
+        if not isinstance(param_states, dict):
+            raise TypeError(f"states must be a dict, got {param_states}")
+        for k, v in param_states.items():
+            if not isinstance(v, State):
+                raise TypeError(f"states values must be ParamState, got {v}")
+        self.param_states.update(param_states)
+        self.param_states.unique_()
+
+        # wrt
+        self.opt_state = ShortTermState(self.tx.init({k: v.value for k, v in self.param_states.items()}))
+        return self
+
+    def update(self, grads: Dict[Hashable, PyTree]):
+        """Update the model states with the gradients.
+
+        Args:
+          grads: the gradients derived from ``brainstate.augment.grad``.
+        """
+        if self.opt_state is None:
+            raise ValueError("register_trainable_weights must be called before update.")
+
+        import optax  # type: ignore[import-not-found,import-untyped]
+        grads = {k: grads[k] for k in self.param_states.keys()}
+        states = {k: v.value for k, v in self.param_states.items()}
+
+        # compute updates
+        updates, new_opt_state = self.tx.update(grads, self.opt_state.value, states)
+        new_params = optax.apply_updates(states, updates)
+
+        # update model states and optimizer states
+        for k, v in self.param_states.items():
+            v.value = new_params[k]
+        self.opt_state.value = new_opt_state

brainstate/optim/_optax_optimizer_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Flax Authors.
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,4 +11,44 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+# ==============================================================================
 
+from __future__ import annotations
+
+import unittest
+
+import jax
+import optax
+
+import brainstate as bst
+
+
+class TestOptaxOptimizer(unittest.TestCase):
+    def test1(self):
+        class Model(bst.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear1 = bst.nn.Linear(2, 3)
+                self.linear2 = bst.nn.Linear(3, 4)
+
+            def __call__(self, x):
+                return self.linear2(self.linear1(x))
+
+        x = bst.random.randn(1, 2)
+        y = jax.numpy.ones((1, 4))
+
+        model = Model()
+        tx = optax.adam(1e-3)
+        optimizer = bst.optim.OptaxOptimizer(tx)
+        optimizer.register_trainable_weights(model.states(bst.ParamState))
+
+        loss_fn = lambda: ((model(x) - y) ** 2).mean()
+        prev_loss = loss_fn()
+
+        grads = bst.augment.grad(loss_fn, model.states(bst.ParamState))()
+        optimizer.update(grads)
+
+        new_loss = loss_fn()
+
+        print(new_loss, prev_loss)
+        self.assertLess(new_loss, prev_loss)