MultiMetric

class brainstate.nn.MultiMetric(**metrics)

Container for multiple metrics updated simultaneously.

This class allows you to group multiple metrics together and update them all with a single call. It’s useful for tracking multiple evaluation metrics (e.g., accuracy, loss, F1 score) during training or evaluation.

Parameters:

**metrics – Keyword arguments where keys are metric names (strings) and values are Metric instances.

_metric_names

List of metric names in the order they were added.

Type:

list of str

Examples

>>> import brainstate
>>> import jax, jax.numpy as jnp
>>> metrics = brainstate.nn.MultiMetric(
...     accuracy=brainstate.nn.AccuracyMetric(),
...     loss=brainstate.nn.AverageMetric(),
... )
>>> logits = jax.random.normal(jax.random.key(0), (5, 2))
>>> labels = jnp.array([1, 1, 0, 1, 0])
>>> batch_loss = jnp.array([1, 2, 3, 4])
>>> 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.reset()
>>> metrics.compute()
{'accuracy': Array(nan, dtype=float32), 'loss': Array(nan, dtype=float32)}

Notes

All keyword arguments passed to update are forwarded to all underlying metrics. Each metric will extract the arguments it needs based on its implementation.

Reserved names (‘reset’, ‘update’, ‘compute’, ‘_metric_names’) cannot be used as metric names.

compute()

Compute and return all metric values.

Returns:

Dictionary mapping metric names to their computed values. The value type depends on the specific metric implementation.

Return type:

dict[str, Any]

Examples

>>> import brainstate
>>> metrics = brainstate.nn.MultiMetric(
...     loss=brainstate.nn.AverageMetric(),
...     stats=brainstate.nn.WelfordMetric(),
... )
>>> # After updates...
>>> results = metrics.compute()
>>> results['loss']  # Returns a scalar
>>> results['stats']  # Returns a Statistics object

reset()

Reset all underlying metrics.

This calls the reset method on each metric in the collection.

Return type:

None

update(**updates)

Update all underlying metrics.

All keyword arguments are passed to the update method of each metric. Individual metrics will extract the arguments they need.

Parameters:

**updates – Keyword arguments to be passed to all underlying metrics.

Return type:

None

Examples

>>> import jax.numpy as jnp
>>> import brainstate
>>> metrics = brainstate.nn.MultiMetric(
...     accuracy=brainstate.nn.AccuracyMetric(),
...     loss=brainstate.nn.AverageMetric('loss_value'),
... )
>>> logits = jnp.array([[0.2, 0.8], [0.9, 0.1]])
>>> labels = jnp.array([1, 0])
>>> loss = jnp.array([0.5, 0.3])
>>> metrics.update(logits=logits, labels=labels, loss_value=loss)