haliax 1.4.dev343__tar.gz → 1.4.dev345__tar.gz

{haliax-1.4.dev343 → haliax-1.4.dev345}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev343
+Version: 1.4.dev345
 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.dev343 → haliax-1.4.dev345}/docs/scan.md

@@ -289,6 +289,12 @@ which is double that required by the default policy, but it reduces the amount o
 Both `save_carries` and `save_inputs` can either be a boolean or the string "offload". If "offload", then the
 checkpointed values will be offloaded to the host during the forward pass, and reloaded during the backward pass.
 
+In addition, you can offload block internals by passing a list of strings to `offload_block_internals`:
+
+```
+policy = ScanCheckpointPolicy(save_carries=True, save_block_internals=["y"], offload_block_internals=["z"])
+```
+
 
 ### Summary of String and Boolean Aliases
 

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

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

{haliax-1.4.dev343 → haliax-1.4.dev345}/src/haliax/__init__.py

@@ -121,13 +121,38 @@ def full_like(a: NamedArray, fill_value: T, dtype: Optional[DTypeLike] = None) -
     return NamedArray(jnp.full_like(a.array, fill_value, dtype=dtype), a.axes)
 
 
-def arange(axis: Axis, *, start=0, step=1, dtype: Optional[DTypeLike] = None) -> NamedArray:
-    """Version of jnp.arange that returns a NamedArray"""
+def arange(axis: AxisSpec, *, start=0, step=1, dtype: Optional[DTypeLike] = None) -> NamedArray:
+    """
+    Version of jnp.arange that returns a NamedArray.
+
+    This version differs from jnp.arange (beyond the obvious NamedArray) in two ways:
+
+    1) It can work with a start that is a tracer (i.e. a JAX expression), whereas jax arange is not able to handle
+    tracers.
+    2) Axis can be more than one axis, in  which case it's equivalent to arange of the product of sizes, followed by
+    reshape.
+
+    Examples
+
+    ```python
+    X, Y = hax.make_axes(X=3, Y=4)
+    # Create a NamedArray along a single axis
+    arr = hax.arange(X)  # equivalent to jnp.arange(0, 3, 1)
+    # 2D
+    arr = hax.arange((X, Y))  # equivalent to jnp.arange(0, 12, 1).reshape(3, 4)
+    ```
+
+    """
+    from haliax.jax_utils import to_jax_shape
+    from haliax.util import ensure_tuple
+
     # if start is a tracer, we need to be a bit cleverer since arange doesn't support tracers
     # return NamedArray(jnp.arange(start, stop, step, dtype=dtype), (axis,))
+    size = axis_size(axis)
 
-    arr = jax.lax.iota(dtype=dtype or jnp.result_type(start), size=axis.size) * step + start
-    return NamedArray(arr, (axis,))
+    arr = jax.lax.iota(dtype=dtype or jnp.result_type(start), size=size) * step + start
+    arr = arr.reshape(to_jax_shape(axis))
+    return NamedArray(arr, ensure_tuple(axis))
 
 
 # TODO: add overrides for arraylike start/stop to linspace, logspace, geomspace

{haliax-1.4.dev343 → haliax-1.4.dev345}/src/haliax/_src/scan.py

@@ -106,6 +106,13 @@ class ScanCheckpointPolicy:
 
     See Also: https://docs.jax.dev/en/latest/gradient-checkpointing.html#custom-policies-for-offload
     """
+
+    offload_block_internals: list[str] = dataclasses.field(default_factory=list)
+    """
+    List of named block internals to offload to the host. This is useful for reducing memory usage on the device
+    while still avoiding rematerialization.
+    """
+
     prevent_cse: bool = False
     """
     Whether to prevent common subexpression elimination in the checkpointed function.
@@ -168,7 +175,6 @@ class ScanCheckpointPolicy:
         if self.disable:
             return callable
         elif self.simple:
-            print("simple")
             return eqx.filter_checkpoint(callable, prevent_cse=self.prevent_cse)
         else:
             policy = self._to_jax_policy(carry_name, input_name)
@@ -202,6 +208,9 @@ class ScanCheckpointPolicy:
         if isinstance(self.save_block_internals, Sequence):
             our_names_to_save.extend(self.save_block_internals)
 
+        if self.offload_block_internals:
+            our_names_to_offload.extend(self.offload_block_internals)
+
         if not our_names_to_save and not our_names_to_offload and not self.save_block_internals:
             return None
 

{haliax-1.4.dev343 → haliax-1.4.dev345}/src/haliax/jax_utils.py

@@ -263,3 +263,10 @@ def multilevel_scan(f, carry, xs, outer_size, length, reverse=False, unroll=1):
             return x
 
     return carry, jax.tree.map(_deshape, scanned)
+
+
+def to_jax_shape(shape):
+    from haliax.core import Axis, ensure_tuple
+
+    shape = ensure_tuple(shape)
+    return tuple(axis.size if isinstance(axis, Axis) else axis for axis in shape)

{haliax-1.4.dev343 → haliax-1.4.dev345}/src/haliax/random.py

@@ -12,7 +12,7 @@ from haliax.core import NamedArray, NamedOrNumeric, broadcast_to
 from haliax.util import ensure_tuple
 
 from .axis import Axis, AxisSelector, AxisSpec, selects_axis
-from .jax_utils import named_call
+from .jax_utils import named_call, to_jax_shape
 from .partitioning import physical_axis_name, physical_axis_size, pspec_for_axis
 
 
@@ -23,7 +23,7 @@ def uniform(
     shape = ensure_tuple(shape)
     minval = broadcast_to(minval, shape).array
     maxval = broadcast_to(maxval, shape).array
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.uniform(key=key, shape=jax_shape, dtype=dtype, minval=minval, maxval=maxval)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
@@ -31,7 +31,7 @@ def uniform(
 @named_call
 def normal(key, shape: AxisSpec, dtype=float):
     shape = ensure_tuple(shape)
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.normal(key=key, shape=jax_shape, dtype=dtype)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
@@ -40,7 +40,7 @@ def normal(key, shape: AxisSpec, dtype=float):
 def bernoulli(key, shape: AxisSpec, p: NamedOrNumeric):
     shape = ensure_tuple(shape)
     p = broadcast_to(p, shape).array
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.bernoulli(key=key, p=p, shape=jax_shape)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
@@ -50,7 +50,7 @@ def randint(key, shape: AxisSpec, minval: NamedOrNumeric, maxval: NamedOrNumeric
     shape = ensure_tuple(shape)
     minval = broadcast_to(minval, shape).array
     maxval = broadcast_to(maxval, shape).array
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.randint(key=key, shape=jax_shape, minval=minval, maxval=maxval, dtype=dtype)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
@@ -59,7 +59,7 @@ def randint(key, shape: AxisSpec, minval: NamedOrNumeric, maxval: NamedOrNumeric
 def poisson(key, shape: AxisSpec, lam: NamedOrNumeric, dtype=int):
     shape = ensure_tuple(shape)
     lam = broadcast_to(lam, shape).array
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.poisson(key=key, lam=lam, shape=jax_shape, dtype=dtype)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
@@ -67,7 +67,7 @@ def poisson(key, shape: AxisSpec, lam: NamedOrNumeric, dtype=int):
 @named_call
 def exponential(key, shape: AxisSpec, dtype=float):
     shape = ensure_tuple(shape)
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.exponential(key=key, shape=jax_shape, dtype=dtype)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
@@ -76,7 +76,7 @@ def exponential(key, shape: AxisSpec, dtype=float):
 def gamma(key, shape: AxisSpec, a: NamedOrNumeric, dtype=float):
     shape = ensure_tuple(shape)
     a = broadcast_to(a, shape).array
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.gamma(key=key, a=a, shape=jax_shape, dtype=dtype)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
@@ -86,7 +86,7 @@ def beta(key, shape: AxisSpec, a: NamedOrNumeric, b: NamedOrNumeric, dtype=float
     shape = ensure_tuple(shape)
     a = broadcast_to(a, shape).array
     b = broadcast_to(b, shape).array
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.beta(key=key, a=a, b=b, shape=jax_shape, dtype=dtype)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
@@ -94,7 +94,7 @@ def beta(key, shape: AxisSpec, a: NamedOrNumeric, b: NamedOrNumeric, dtype=float
 @named_call
 def laplace(key, shape: AxisSpec, dtype=float):
     shape = ensure_tuple(shape)
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.laplace(key=key, shape=jax_shape, dtype=dtype)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
@@ -102,7 +102,7 @@ def laplace(key, shape: AxisSpec, dtype=float):
 @named_call
 def cauchy(key, shape: AxisSpec, dtype=float):
     shape = ensure_tuple(shape)
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.cauchy(key=key, shape=jax_shape, dtype=dtype)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
@@ -110,7 +110,7 @@ def cauchy(key, shape: AxisSpec, dtype=float):
 @named_call
 def logistic(key, shape: AxisSpec, dtype=float):
     shape = ensure_tuple(shape)
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.logistic(key=key, shape=jax_shape, dtype=dtype)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
@@ -120,7 +120,7 @@ def truncated_normal(key, shape: AxisSpec, lower: NamedOrNumeric, upper: NamedOr
     shape = ensure_tuple(shape)
     lower = broadcast_to(lower, shape).array
     upper = broadcast_to(upper, shape).array
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.truncated_normal(key=key, lower=lower, upper=upper, shape=jax_shape, dtype=dtype)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
@@ -201,7 +201,7 @@ def generate_sharded(fn, axis: Optional[AxisSelector] = None):
 @named_call
 def ball(key, shape: AxisSpec, D: Axis, p: float = 2.0, dtype=float):
     shape = ensure_tuple(shape)
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.ball(key=key, shape=jax_shape, d=D.size, p=p, dtype=dtype)
     return haliax.auto_sharded(NamedArray(jax_array, shape + (D,)))
 
@@ -225,7 +225,7 @@ def choice(
     if p is not None:
         assert p.resolve_axis(ensure_tuple(axis)) == p.axes, f"p must be 1D with axis {axis} or be None"
 
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_p = p.array if p is not None else None
 
     jax_array = jrandom.choice(key, a.array, jax_shape, replace=replace, p=jax_p, axis=index)
@@ -264,7 +264,7 @@ def categorical(key, logits: NamedArray, axis: AxisSelector, shape: Optional[Axi
     index = logits._lookup_indices(axis)
     assert index is not None, f"axis {axis} not in logits"
 
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
 
     jax_array = jrandom.categorical(key, logits.array, axis=index, shape=jax_shape)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
@@ -273,7 +273,7 @@ def categorical(key, logits: NamedArray, axis: AxisSelector, shape: Optional[Axi
 @named_call
 def gumbel(key, shape: AxisSpec, dtype=float):
     shape = ensure_tuple(shape)
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.gumbel(key, jax_shape, dtype=dtype)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
@@ -288,7 +288,7 @@ def permutation(key, x: NamedArray, axis: AxisSelector, independent: bool = Fals
 @named_call
 def rademacher(key, shape: AxisSpec, dtype=float):
     shape = ensure_tuple(shape)
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.rademacher(key, jax_shape, dtype=dtype)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
@@ -297,7 +297,7 @@ def rademacher(key, shape: AxisSpec, dtype=float):
 def t(key, shape: AxisSpec, df: NamedOrNumeric, dtype=float):
     shape = ensure_tuple(shape)
     df = broadcast_to(df, shape)
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.t(key, df.array, jax_shape, dtype=dtype)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
@@ -307,7 +307,7 @@ def weibull_min(key, shape: AxisSpec, scale: NamedOrNumeric, concentration: Name
     shape = ensure_tuple(shape)
     scale = broadcast_to(scale, shape)
     concentration = broadcast_to(concentration, shape)
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.weibull_min(key, scale.array, concentration.array, jax_shape, dtype=dtype)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
@@ -316,7 +316,7 @@ def weibull_min(key, shape: AxisSpec, scale: NamedOrNumeric, concentration: Name
 def pareto(key, shape: AxisSpec, b: NamedOrNumeric, dtype=float):
     shape = ensure_tuple(shape)
     b = broadcast_to(b, shape)
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.pareto(key, b.array, jax_shape, dtype=dtype)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
@@ -325,15 +325,11 @@ def pareto(key, shape: AxisSpec, b: NamedOrNumeric, dtype=float):
 def loggamma(key, shape: AxisSpec, a: NamedOrNumeric, dtype=float):
     shape = ensure_tuple(shape)
     a = broadcast_to(a, shape)
-    jax_shape = _to_jax_shape(shape)
+    jax_shape = to_jax_shape(shape)
     jax_array = jrandom.loggamma(key, a.array, jax_shape, dtype=dtype)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 
 
-def _to_jax_shape(shape):
-    return tuple(axis.size if isinstance(axis, Axis) else axis for axis in shape)
-
-
 __all__ = [
     "generate_sharded",
     "uniform",

{haliax-1.4.dev343 → haliax-1.4.dev345}/tests/test_scan.py

@@ -171,40 +171,56 @@ E = hax.Axis("E", 10)
 
 
 @pytest.mark.parametrize(
-    "name,policy,expected_scan_shapes",
+    "name,policy,expected_scan_shapes,check_offloading",
     [
-        ("disabled", ScanCheckpointPolicy(disable=True), [(E.size,), (Block.size, E.size), (Block.size, E.size)]),
-        ("carry_true", True, [(E.size,), (Block.size, E.size)]),
+        (
+            "disabled",
+            ScanCheckpointPolicy(disable=True),
+            [(E.size,), (Block.size, E.size), (Block.size, E.size)],
+            None,
+        ),
+        ("carry_true", True, [(E.size,), (Block.size, E.size)], None),
         (
             "carry",
             ScanCheckpointPolicy(save_carries=True, save_block_internals=False),
             [(E.size,), (Block.size, E.size)],
+            None,
         ),
         (
             "everything",
             ScanCheckpointPolicy(save_carries=True, save_inputs=True, save_block_internals=True),
             [(E.size,), (Block.size, E.size), (Block.size, E.size)],
+            None,
         ),
         (
             "internals",
             ScanCheckpointPolicy(save_carries=False, save_block_internals=True),
             [(E.size,), (Block.size, E.size), (Block.size, E.size)],
+            None,
         ),
         (
             "cos",
             ScanCheckpointPolicy(save_carries=False, save_block_internals=["cos"]),
             [(E.size,), (Block.size, E.size)],
+            None,
         ),
         (
             "sin",
             ScanCheckpointPolicy(save_carries=True, save_block_internals=["sin"]),
             [(E.size,), (Block.size, E.size), (Block.size, E.size)],
+            None,
+        ),
+        ("simple", ScanCheckpointPolicy(simple=True), [(E.size,), (Block.size, E.size)], None),
+        ("nested", ScanCheckpointPolicy(simple=True, nested=2), [(E.size,), (2, E.size)], None),
+        (
+            "sin_offload",
+            ScanCheckpointPolicy(save_carries=True, offload_block_internals=["sin"]),
+            [(E.size,), (Block.size, E.size), (Block.size, E.size)],
+            ["sin"],
         ),
-        ("simple", ScanCheckpointPolicy(simple=True), [(E.size,), (Block.size, E.size)]),
-        ("nested", ScanCheckpointPolicy(simple=True, nested=2), [(E.size,), (2, E.size)]),
     ],
 )
-def test_checkpoint_carries(name, policy, expected_scan_shapes):
+def test_checkpoint_carries(name, policy, expected_scan_shapes, check_offloading):
     class Module(eqx.Module):
         named: hax.NamedArray
 
@@ -243,3 +259,29 @@ def test_checkpoint_carries(name, policy, expected_scan_shapes):
         print(residual)
 
     assert out_shapes == expected_scan_shapes, f"{name}: Expected {expected_scan_shapes}, got {out_shapes}"
+
+    # Add check for offloading if specified
+    if check_offloading is not None:
+        for name in check_offloading:
+            print(f"Checking offloading for {name}")
+            target = None
+            found_saved = False
+            for expr in jaxpr.jaxpr.eqns:
+                if expr.primitive.name == "scan":
+                    inner_jaxpr = expr.params["jaxpr"]
+                    for eqn in inner_jaxpr.eqns:
+                        if eqn.primitive.name == "name":
+                            this_name = eqn.params["name"]
+                            if this_name == name:
+                                # TODO in theory we can save more than one thing with the same name
+                                # not gonna worry about that for now
+                                target = eqn.outvars[0]
+                        elif eqn.primitive.name == "device_put":
+                            if eqn.invars[0] == target:
+                                found_saved = True
+                                break
+                    # found scan
+                    break
+
+            assert target is not None, f"Could not find named value for {name}"
+            assert found_saved, f"Could not find offloaded value for {name}"

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

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