brainstate 0.1.10__py2.py3-none-any.whl → 0.2.0__py2.py3-none-any.whl

brainstate/nn/metrics.py

@@ -1,388 +0,0 @@
-# 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 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 brainstate
-
-      >>> batch_loss = jnp.array([1, 2, 3, 4])
-      >>> batch_loss2 = jnp.array([3, 2, 1, 0])
-
-      >>> metrics = brainstate.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 brainstate
-      >>> 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 = brainstate.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

@@ -1,38 +0,0 @@
-# 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 ._base import *
-from ._base import __all__ as base_all
-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__ = (
-    base_all
-    + scheduler_all
-    + optimizer_all
-    + optax_all
-)
-
-del (
-    base_all,
-    optax_all,
-    scheduler_all,
-    optimizer_all,
-)

brainstate/optim/_base.py

@@ -1,64 +0,0 @@
-# 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 typing import Dict, Hashable
-
-from brainstate._state import State, StateDictManager
-from brainstate.graph import Node
-from brainstate.typing import PyTree
-
-__all__ = [
-    'Optimizer',
-]
-
-
-class Optimizer(Node):
-    """
-    Base Optimizer Class.
-
-    Attributes:
-    -----------
-    param_states: StateDictManager
-        The state dict manager for trainable weights.
-
-    """
-
-    param_states: StateDictManager
-
-    def __init__(self):
-        self.param_states = StateDictManager()
-
-    def register_trainable_weights(self, param_states: Dict[Hashable, State]):
-        """
-        Register the trainable weights.
-
-        Parameters:
-        -----------
-        param_states: Dict[Hashable, State]
-            The trainable weights.
-        """
-        raise NotImplementedError
-
-    def update(self, grads: Dict[Hashable, PyTree]):
-        """
-        Update the trainable weights according to weight gradients.
-
-        Parameters:
-        -----------
-        grads: Dict[Hashable, PyTree]
-            The weight gradients.
-        """
-        raise NotImplementedError