brainstate 0.0.2.post20240913__py2.py3-none-any.whl → 0.0.2.post20241009__py2.py3-none-any.whl

brainstate/nn/event/linear_test.py

@@ -0,0 +1,111 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+# ==============================================================================
+
+
+import jax
+import jax.numpy as jnp
+from absl.testing import parameterized
+
+import brainstate as bst
+from brainstate.nn.event.linear import EventDense
+
+
+class TestEventLinear(parameterized.TestCase):
+  @parameterized.product(
+    homo_w=[True, False],
+    bool_x=[True, False],
+  )
+  def test1(self, homo_w, bool_x):
+    x = bst.random.rand(20) < 0.1
+    if not bool_x:
+      x = jnp.asarray(x, dtype=float)
+    m = EventDense(20, 40, 1.5 if homo_w else bst.init.KaimingUniform())
+    y = m(x)
+    print(y)
+
+    self.assertTrue(jnp.allclose(y, (x.sum() * m.weight) if homo_w else (x @ m.weight)))
+
+  def test_grad_bool(self):
+    n_in = 20
+    n_out = 30
+    x = bst.random.rand(n_in) < 0.3
+    fn = EventDense(n_in, n_out, bst.init.KaimingUniform())
+
+    with self.assertRaises(TypeError):
+      print(jax.grad(lambda x: fn(x).sum())(x))
+
+  @parameterized.product(
+    bool_x=[True, False],
+    homo_w=[True, False]
+  )
+  def test_vjp(self, bool_x, homo_w):
+    n_in = 20
+    n_out = 30
+    if bool_x:
+      x = jax.numpy.asarray(bst.random.rand(n_in) < 0.3, dtype=float)
+    else:
+      x = bst.random.rand(n_in)
+
+    fn = EventDense(n_in, n_out, 1.5 if homo_w else bst.init.KaimingUniform())
+    w = fn.weight
+
+    def f(x, w):
+      fn.weight = w
+      return fn(x).sum()
+
+    r1 = jax.grad(f, argnums=(0, 1))(x, w)
+
+    # -------------------
+    # TRUE gradients
+
+    def f2(x, w):
+      y = (x @ (jnp.ones([n_in, n_out]) * w)) if homo_w else (x @ w)
+      return y.sum()
+
+    r2 = jax.grad(f2, argnums=(0, 1))(x, w)
+    self.assertTrue(jnp.allclose(r1[0], r2[0]))
+    self.assertTrue(jnp.allclose(r1[1], r2[1]))
+
+  @parameterized.product(
+    bool_x=[True, False],
+    homo_w=[True, False]
+  )
+  def test_jvp(self, bool_x, homo_w):
+    n_in = 20
+    n_out = 30
+    if bool_x:
+      x = jax.numpy.asarray(bst.random.rand(n_in) < 0.3, dtype=float)
+    else:
+      x = bst.random.rand(n_in)
+
+    fn = EventDense(n_in, n_out, 1.5 if homo_w else bst.init.KaimingUniform(), grad_mode='jvp')
+    w = fn.weight
+
+    def f(x, w):
+      fn.weight = w
+      return fn(x)
+
+    o1, r1 = jax.jvp(f, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
+
+    # -------------------
+    # TRUE gradients
+
+    def f2(x, w):
+      y = (x @ (jnp.ones([n_in, n_out]) * w)) if homo_w else (x @ w)
+      return y
+
+    o2, r2 = jax.jvp(f, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
+    self.assertTrue(jnp.allclose(o1, o2))
+    self.assertTrue(jnp.allclose(r1, r2))

brainstate/nn/metrics.py

@@ -0,0 +1,390 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 typing as tp
+from dataclasses import dataclass
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from brainstate._state import State
+
+__all__ = [
+  'Average',
+  'Statistics',
+  'Welford',
+  'Accuracy',
+  'MultiMetric',
+]
+
+
+class MetricState(State):
+  """Wrapper class for Metric Variables."""
+  pass
+
+
+class Metric(object):
+  """Base class for metrics. Any class that subclasses ``Metric`` should
+  implement a ``compute``, ``reset`` and ``update`` method."""
+
+  def reset(self) -> None:
+    """In-place reset the ``Metric``."""
+    raise NotImplementedError('Must override `reset()` method.')
+
+  def update(self, **kwargs) -> None:
+    """In-place update the ``Metric``."""
+    raise NotImplementedError('Must override `update()` method.')
+
+  def compute(self):
+    """Compute and return the value of the ``Metric``."""
+    raise NotImplementedError('Must override `compute()` method.')
+
+
+class Average(Metric):
+  """Average metric.
+
+  Example usage::
+
+    >>> import jax.numpy as jnp
+    >>> import brainstate as bst
+
+    >>> batch_loss = jnp.array([1, 2, 3, 4])
+    >>> batch_loss2 = jnp.array([3, 2, 1, 0])
+
+    >>> metrics = bst.nn.metrics.Average()
+    >>> metrics.compute()
+    Array(nan, dtype=float32)
+    >>> metrics.update(values=batch_loss)
+    >>> metrics.compute()
+    Array(2.5, dtype=float32)
+    >>> metrics.update(values=batch_loss2)
+    >>> metrics.compute()
+    Array(2., dtype=float32)
+    >>> metrics.reset()
+    >>> metrics.compute()
+    Array(nan, dtype=float32)
+  """
+
+  def __init__(self, argname: str = 'values'):
+    """Pass in a string denoting the key-word argument that :func:`update` will use to derive the new value.
+    For example, constructing the metric as ``avg = Average('test')`` would allow you to make updates with
+    ``avg.update(test=new_value)``.
+
+    Args:
+      argname: an optional string denoting the key-word argument that
+        :func:`update` will use to derive the new value. Defaults to
+        ``'values'``.
+    """
+    self.argname = argname
+    self.total = MetricState(jnp.array(0, dtype=jnp.float32))
+    self.count = MetricState(jnp.array(0, dtype=jnp.int32))
+
+  def reset(self) -> None:
+    """Reset this ``Metric``."""
+    self.total.value = jnp.array(0, dtype=jnp.float32)
+    self.count.value = jnp.array(0, dtype=jnp.int32)
+
+  def update(self, **kwargs) -> None:
+    """In-place update this ``Metric``. This method will use the value from
+    ``kwargs[self.argname]`` to update the metric, where ``self.argname`` is
+    defined on construction.
+
+    Args:
+      **kwargs: the key-word arguments that contains a ``self.argname``
+        entry that maps to the value we want to use to update this metric.
+    """
+    if self.argname not in kwargs:
+      raise TypeError(f"Expected keyword argument '{self.argname}'")
+    values: tp.Union[int, float, jax.Array] = kwargs[self.argname]
+    self.total.value += (
+      values if isinstance(values, (int, float)) else values.sum()
+    )
+    self.count.value += 1 if isinstance(values, (int, float)) else values.size
+
+  def compute(self) -> jax.Array:
+    """Compute and return the average."""
+    return self.total.value / self.count.value
+
+
+@partial(jax.tree_util.register_dataclass,
+         data_fields=['mean', 'standard_error_of_mean', 'standard_deviation'],
+         meta_fields=[])
+@dataclass
+class Statistics:
+  mean: jnp.float32
+  standard_error_of_mean: jnp.float32
+  standard_deviation: jnp.float32
+
+
+class Welford(Metric):
+  """Uses Welford's algorithm to compute the mean and variance of a stream of data.
+
+  Example usage::
+
+    >>> import jax.numpy as jnp
+    >>> from brainstate import nn
+
+    >>> batch_loss = jnp.array([1, 2, 3, 4])
+    >>> batch_loss2 = jnp.array([3, 2, 1, 0])
+
+    >>> metrics = nn.metrics.Welford()
+    >>> metrics.compute()
+    Statistics(mean=Array(0., dtype=float32), standard_error_of_mean=Array(nan, dtype=float32), standard_deviation=Array(nan, dtype=float32))
+    >>> metrics.update(values=batch_loss)
+    >>> metrics.compute()
+    Statistics(mean=Array(2.5, dtype=float32), standard_error_of_mean=Array(0.559017, dtype=float32), standard_deviation=Array(1.118034, dtype=float32))
+    >>> metrics.update(values=batch_loss2)
+    >>> metrics.compute()
+    Statistics(mean=Array(2., dtype=float32), standard_error_of_mean=Array(0.43301272, dtype=float32), standard_deviation=Array(1.2247449, dtype=float32))
+    >>> metrics.reset()
+    >>> metrics.compute()
+    Statistics(mean=Array(0., dtype=float32), standard_error_of_mean=Array(nan, dtype=float32), standard_deviation=Array(nan, dtype=float32))
+  """
+
+  def __init__(self, argname: str = 'values'):
+    """Pass in a string denoting the key-word argument that :func:`update` will use to derive the new value.
+    For example, constructing the metric as ``wf = Welford('test')`` would allow you to make updates with
+    ``wf.update(test=new_value)``.
+
+    Args:
+      argname: an optional string denoting the key-word argument that
+        :func:`update` will use to derive the new value. Defaults to
+        ``'values'``.
+    """
+    self.argname = argname
+    self.count = MetricState(jnp.array(0, dtype=jnp.int32))
+    self.mean = MetricState(jnp.array(0, dtype=jnp.float32))
+    self.m2 = MetricState(jnp.array(0, dtype=jnp.float32))
+
+  def reset(self) -> None:
+    """Reset this ``Metric``."""
+    self.count.value = jnp.array(0, dtype=jnp.uint32)
+    self.mean.value = jnp.array(0, dtype=jnp.float32)
+    self.m2.value = jnp.array(0, dtype=jnp.float32)
+
+  def update(self, **kwargs) -> None:
+    """In-place update this ``Metric``. This method will use the value from
+    ``kwargs[self.argname]`` to update the metric, where ``self.argname`` is
+    defined on construction.
+
+    Args:
+      **kwargs: the key-word arguments that contains a ``self.argname``
+        entry that maps to the value we want to use to update this metric.
+    """
+    if self.argname not in kwargs:
+      raise TypeError(f"Expected keyword argument '{self.argname}'")
+    values: tp.Union[int, float, jax.Array] = kwargs[self.argname]
+    count = 1 if isinstance(values, (int, float)) else values.size
+    original_count = self.count.value
+    self.count.value += count
+    delta = (
+              values if isinstance(values, (int, float)) else values.mean()
+            ) - self.mean.value
+    self.mean.value += delta * count / self.count.value
+    m2 = 0.0 if isinstance(values, (int, float)) else values.var() * count
+    self.m2.value += (
+        m2 + delta * delta * count * original_count / self.count
+    )
+
+  def compute(self) -> Statistics:
+    """Compute and return the mean and variance statistics in a
+    ``Statistics`` dataclass object.
+    """
+    variance = self.m2.value / self.count.value
+    standard_deviation = variance ** 0.5
+    sem = standard_deviation / (self.count.value ** 0.5)
+    return Statistics(
+      mean=self.mean.value,
+      standard_error_of_mean=sem,
+      standard_deviation=standard_deviation,
+    )
+
+
+class Accuracy(Average):
+  """Accuracy metric. This metric subclasses :class:`Average`,
+  and so they share the same ``reset`` and ``compute`` method
+  implementations. Unlike :class:`Average`, no string needs to
+  be passed to ``Accuracy`` during construction.
+
+  Example usage::
+
+    >>> import brainstate as bst
+    >>> import jax, jax.numpy as jnp
+
+    >>> logits = jax.random.normal(jax.random.key(0), (5, 2))
+    >>> labels = jnp.array([1, 1, 0, 1, 0])
+    >>> logits2 = jax.random.normal(jax.random.key(1), (5, 2))
+    >>> labels2 = jnp.array([0, 1, 1, 1, 1])
+
+    >>> metrics = bst.nn.metrics.Accuracy()
+    >>> metrics.compute()
+    Array(nan, dtype=float32)
+    >>> metrics.update(logits=logits, labels=labels)
+    >>> metrics.compute()
+    Array(0.6, dtype=float32)
+    >>> metrics.update(logits=logits2, labels=labels2)
+    >>> metrics.compute()
+    Array(0.7, dtype=float32)
+    >>> metrics.reset()
+    >>> metrics.compute()
+    Array(nan, dtype=float32)
+  """
+
+  def update(self, *, logits: jax.Array, labels: jax.Array, **_) -> None:  # type: ignore[override]
+    """In-place update this ``Metric``.
+
+    Args:
+      logits: the outputted predicted activations. These values are
+        argmax-ed (on the trailing dimension), before comparing them
+        to the labels.
+      labels: the ground truth integer labels.
+    """
+    if logits.ndim != labels.ndim + 1:
+      raise ValueError(
+        f'Expected logits.ndim==labels.ndim+1, got {logits.ndim} and {labels.ndim}'
+      )
+    elif labels.dtype in (jnp.int64, np.int32, np.int64):
+      labels = jnp.astype(labels, jnp.int32)
+    elif labels.dtype != jnp.int32:
+      raise ValueError(f'Expected labels.dtype==jnp.int32, got {labels.dtype}')
+
+    super().update(values=(logits.argmax(axis=-1) == labels))
+
+
+class MultiMetric(Metric):
+  """MultiMetric class to store multiple metrics and update them in a single call.
+
+  Example usage::
+
+    >>> from brainstate import nn
+    >>> import jax, jax.numpy as jnp
+
+    >>> metrics = nn.metrics.MultiMetric(
+    ...   accuracy=nn.metrics.Accuracy(),
+    ...   loss=nn.metrics.Average(),
+    ... )
+
+    >>> metrics
+    MultiMetric(
+      accuracy=Accuracy(
+        argname='values',
+        total=MetricState(
+          value=Array(0., dtype=float32)
+        ),
+        count=MetricState(
+          value=Array(0, dtype=int32)
+        )
+      ),
+      loss=Average(
+        argname='values',
+        total=MetricState(
+          value=Array(0., dtype=float32)
+        ),
+        count=MetricState(
+          value=Array(0, dtype=int32)
+        )
+      )
+    )
+
+    >>> metrics.accuracy
+    Accuracy(
+      argname='values',
+      total=MetricState(
+        value=Array(0., dtype=float32)
+      ),
+      count=MetricState(
+        value=Array(0, dtype=int32)
+      )
+    )
+
+    >>> metrics.loss
+    Average(
+      argname='values',
+      total=MetricState(
+        value=Array(0., dtype=float32)
+      ),
+      count=MetricState(
+        value=Array(0, dtype=int32)
+      )
+    )
+
+    >>> logits = jax.random.normal(jax.random.key(0), (5, 2))
+    >>> labels = jnp.array([1, 1, 0, 1, 0])
+    >>> logits2 = jax.random.normal(jax.random.key(1), (5, 2))
+    >>> labels2 = jnp.array([0, 1, 1, 1, 1])
+
+    >>> batch_loss = jnp.array([1, 2, 3, 4])
+    >>> batch_loss2 = jnp.array([3, 2, 1, 0])
+
+    >>> metrics.compute()
+    {'accuracy': Array(nan, dtype=float32), 'loss': Array(nan, dtype=float32)}
+    >>> metrics.update(logits=logits, labels=labels, values=batch_loss)
+    >>> metrics.compute()
+    {'accuracy': Array(0.6, dtype=float32), 'loss': Array(2.5, dtype=float32)}
+    >>> metrics.update(logits=logits2, labels=labels2, values=batch_loss2)
+    >>> metrics.compute()
+    {'accuracy': Array(0.7, dtype=float32), 'loss': Array(2., dtype=float32)}
+    >>> metrics.reset()
+    >>> metrics.compute()
+    {'accuracy': Array(nan, dtype=float32), 'loss': Array(nan, dtype=float32)}
+  """
+
+  def __init__(self, **metrics):
+    """Pass in key-word arguments to the constructor, e.g.
+    ``MultiMetric(keyword1=Average(), keyword2=Accuracy(), ...)``.
+
+    Args:
+      **metrics: the key-word arguments that will be used to access
+        the corresponding ``Metric``.
+    """
+    # TODO: raise error if a kwarg is passed that is in ('reset', 'update', 'compute'), since these names are reserved for methods
+    self._metric_names = []
+    for metric_name, metric in metrics.items():
+      self._metric_names.append(metric_name)
+      vars(self)[metric_name] = metric
+
+  def reset(self) -> None:
+    """Reset all underlying ``Metric``'s."""
+    for metric_name in self._metric_names:
+      getattr(self, metric_name).reset()
+
+  def update(self, **updates) -> None:
+    """In-place update all underlying ``Metric``'s in this ``MultiMetric``. All
+    ``**updates`` will be passed to the ``update`` method of all underlying
+    ``Metric``'s.
+
+    Args:
+      **updates: the key-word arguments that will be passed to the underlying ``Metric``'s
+        ``update`` method.
+    """
+    # TODO: should we give the option of updating only some of the metrics and not all? e.g. if for some kwargs==None, don't do update
+    # TODO: should we raise an error if a kwarg is passed into **updates that has no match with any underlying metric? e.g. user typo
+    for metric_name in self._metric_names:
+      getattr(self, metric_name).update(**updates)
+
+  def compute(self) -> dict[str, Metric]:
+    """Compute and return the value of all underlying ``Metric``'s. This method
+    will return a dictionary, mapping strings (defined by the key-word arguments
+    ``**metrics`` passed to the constructor) to the corresponding metric value.
+    """
+    return {
+      f'{metric_name}': getattr(self, metric_name).compute()
+      for metric_name in self._metric_names
+    }

brainstate/optim/__init__.py

@@ -16,7 +16,11 @@
 
 from ._lr_scheduler import *
 from ._lr_scheduler import __all__ as scheduler_all
+from ._optax_optimizer import *
+from ._optax_optimizer import __all__ as optax_all
 from ._sgd_optimizer import *
 from ._sgd_optimizer import __all__ as optimizer_all
 
-__all__ = scheduler_all + optimizer_all
+__all__ = scheduler_all + optimizer_all + optax_all
+
+del optax_all, scheduler_all, optimizer_all