haliax 1.4.dev306__tar.gz → 1.4.dev308__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: haliax
-Version: 1.4.dev306
+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.dev306 → 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.dev306 → 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.dev306 → haliax-1.4.dev308}/src/haliax/partitioning.py

@@ -139,10 +139,11 @@ def shard(x: T, mapping: Optional[ResourceMapping] = None, mesh: Optional[Mesh]
 
         sharding = infer_resource_partitions(named, mapping, mesh=mesh, preserve_existing_shardings=False)
         assert isinstance(sharding, NamedSharding)
+        in_sharding = getattr(named.array, "sharding", None)
         if is_in_jit():
             return with_sharding_constraint(named, sharding)
         # as a special case, SingleDeviceShardings are routed through jit
-        elif isinstance(named.array.sharding, SingleDeviceSharding):
+        elif isinstance(in_sharding, SingleDeviceSharding) and in_sharding._device in sharding.device_set:
             # TODO(dlwh): this should be unnecessary in JAX soon. Check after 2024-08-01
             sharded_array = jax.jit(lambda x: x, out_shardings=sharding)(named)
             return sharded_array

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

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