haliax 1.4.dev307__tar.gz → 1.4.dev308__tar.gz

{haliax-1.4.dev307 → haliax-1.4.dev308}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: haliax
-Version: 1.4.dev307
+Version: 1.4.dev308
 Summary: Named Tensors for Legible Deep Learning in JAX
 Project-URL: Homepage, https://github.com/stanford-crfm/haliax
 Project-URL: Bug Tracker, https://github.com/stanford-crfm/haliax/issues/

haliax-1.4.dev308/src/haliax/__about__.py

@@ -0,0 +1 @@
+__version__ = "1.4.dev308"

{haliax-1.4.dev307 → haliax-1.4.dev308}/src/haliax/nn/__init__.py

@@ -34,7 +34,7 @@ from .conv import Conv, ConvTranspose
 from .dropout import Dropout, dropout
 from .embedding import Embedding
 from .linear import Linear
-from .loss import binary_cross_entropy_loss, cross_entropy_loss, cross_entropy_loss_and_log_normalizers
+from .loss import binary_cross_entropy_loss, cross_entropy_loss, cross_entropy_loss_and_log_normalizers, reduce_loss
 from .mlp import MLP
 from .normalization import LayerNorm, log_softmax, logsumexp, softmax, standardize
 from .pool import max_pool, mean_pool, min_pool
@@ -77,6 +77,7 @@ __all__ = [
     "attention",
     "one_hot",
     "binary_cross_entropy_loss",
+    "reduce_loss",
     "cross_entropy_loss",
     "cross_entropy_loss_and_log_normalizers",
     "Conv",

{haliax-1.4.dev307 → haliax-1.4.dev308}/src/haliax/nn/loss.py

@@ -94,6 +94,21 @@ def binary_cross_entropy_loss(
     return loss
 
 
+def reduce_loss(
+    arr,
+    reduction: Optional[ReductionFunction] | Unspecified = UNSPECIFIED,
+    reduction_axis: Optional[AxisSelection] = None,
+    where: Optional[NamedArray] = None,
+):
+    """
+    Reduce a loss array according to the given reduction and reduction axis.
+    If reduction is None, the loss is not reduced.
+    If reduction is UNSPECIFIED, the default reduction is used (mean).
+    If reduction_axis is None (default), the loss is reduced over all axes.
+    """
+    return maybe_reduce_loss(arr, reduction, reduction_axis, where)
+
+
 def maybe_reduce_loss(
     arr,
     reduction: Optional[ReductionFunction] | Unspecified,

haliax-1.4.dev307/src/haliax/__about__.py

@@ -1 +0,0 @@
-__version__ = "1.4.dev307"