brainstate 0.0.2.post20241010__py2.py3-none-any.whl → 0.1.0.post20241122__py2.py3-none-any.whl

brainstate/nn/metrics.py

@@ -27,100 +27,100 @@ import numpy as np
 from brainstate._state import State
 
 __all__ = [
-  'Average',
-  'Statistics',
-  'Welford',
-  'Accuracy',
-  'MultiMetric',
+    'Average',
+    'Statistics',
+    'Welford',
+    'Accuracy',
+    'MultiMetric',
 ]
 
 
 class MetricState(State):
-  """Wrapper class for Metric Variables."""
-  pass
+    """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."""
+    """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 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 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.')
+    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'``.
+    """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)
     """
-    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
+    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,
@@ -128,162 +128,183 @@ class Average(Metric):
          meta_fields=[])
 @dataclass
 class Statistics:
-  mean: jnp.float32
-  standard_error_of_mean: jnp.float32
-  standard_deviation: jnp.float32
+    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'``.
+    """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))
     """
-    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,
-    )
+
+    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.
+    """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)
     """
-    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))
 
+    def update(self, *, logits: jax.Array, labels: jax.Array, **_) -> None:  # type: ignore[override]
+        """In-place update this ``Metric``.
 
-class MultiMetric(Metric):
-  """MultiMetric class to store multiple metrics and update them in a single call.
+        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}')
 
-  Example usage::
+        super().update(values=(logits.argmax(axis=-1) == labels))
 
-    >>> from brainstate import nn
-    >>> import jax, jax.numpy as jnp
 
-    >>> metrics = nn.metrics.MultiMetric(
-    ...   accuracy=nn.metrics.Accuracy(),
-    ...   loss=nn.metrics.Average(),
-    ... )
+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
-    MultiMetric(
-      accuracy=Accuracy(
+      >>> metrics.accuracy
+      Accuracy(
         argname='values',
         total=MetricState(
           value=Array(0., dtype=float32)
@@ -291,8 +312,10 @@ class MultiMetric(Metric):
         count=MetricState(
           value=Array(0, dtype=int32)
         )
-      ),
-      loss=Average(
+      )
+
+      >>> metrics.loss
+      Average(
         argname='values',
         total=MetricState(
           value=Array(0., dtype=float32)
@@ -301,90 +324,67 @@ class MultiMetric(Metric):
           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.
+
+      >>> 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)}
     """
-    return {
-      f'{metric_name}': getattr(self, metric_name).compute()
-      for metric_name in self._metric_names
-    }
+
+    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
+        }