haliax 1.4.dev296__tar.gz → 1.4.dev298__tar.gz

{haliax-1.4.dev296 → haliax-1.4.dev298}/.github/workflows/run_tests.yaml

@@ -21,13 +21,7 @@ jobs:
           python -m pip install --upgrade pip
           pip install flake8 pytest
           pip install --upgrade "jax[cpu]==0.4.25" "jaxlib[cpu]==0.4.25"
-          pip install .
-#      - name: Lint with flake8
-#        run: |
-#          # stop the build if there are Python syntax errors or undefined names
-#          flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
-#          # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
-#          flake8 . --count --exit-zero --max-complexity=50 --max-line-length=127 --statistics
+          pip install .[dev]
       - name: Test with pytest
         run: |
           XLA_FLAGS=--xla_force_host_platform_device_count=8 PYTHONPATH=tests:src:. pytest tests

{haliax-1.4.dev296 → haliax-1.4.dev298}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: haliax
-Version: 1.4.dev296
+Version: 1.4.dev298
 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/
@@ -17,6 +17,7 @@ Requires-Dist: equinox>=0.10.6
 Requires-Dist: jaxtyping>=0.2.20
 Requires-Dist: jmp>=0.0.4
 Provides-Extra: dev
+Requires-Dist: chex>=0.1.86; extra == 'dev'
 Requires-Dist: mkdocs-include-markdown-plugin; extra == 'dev'
 Requires-Dist: mkdocs-literate-nav>=0.6.0; extra == 'dev'
 Requires-Dist: mkdocs-macros-plugin>=0.7.0; extra == 'dev'

{haliax-1.4.dev296 → haliax-1.4.dev298}/pyproject.toml

@@ -36,8 +36,10 @@ dev=["pytest >= 7.4.0", "mypy >= 0.910", "mkdocs >= 1.4.3", "mkdocs-material >=
     "pymdown-extensions",
     "pygments",
     "pymdown-extensions",
+    "chex>=0.1.86"
 ]
 
+
 [tool.hatch.version]
 path = "src/haliax/__about__.py"
 

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

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

{haliax-1.4.dev296 → haliax-1.4.dev298}/src/haliax/__init__.py

@@ -26,6 +26,7 @@ from .axis import (
     AxisSelector,
     AxisSpec,
     axis_name,
+    axis_size,
     concat_axes,
     dblock,
     ds,
@@ -55,6 +56,7 @@ from .core import (
     unflatten_axis,
     updated_slice,
 )
+from .jax_utils import filter_checkpoint
 from .hof import fold, map, scan, vmap
 from .ops import clip, isclose, pad_left, trace, tril, triu, where
 from .partitioning import auto_sharded, axis_mapping, fsdp, named_jit, shard, shard_with_axis_mapping

{haliax-1.4.dev296 → haliax-1.4.dev298}/src/haliax/axis.py

@@ -1,5 +1,6 @@
 import typing
 from dataclasses import dataclass
+from math import prod
 from types import EllipsisType
 from typing import Dict, List, Mapping, Optional, Sequence, Tuple, Union, overload
 
@@ -354,6 +355,17 @@ def axis_name(ax: AxisSelection) -> Union[str, Tuple[str, ...]]:
         return tuple(_ax_name(x) for x in ax)
 
 
+def axis_size(ax: AxisSpec) -> int:
+    """
+    Returns the size of the axis or the product of the sizes of the axes in the axis spec
+    """
+
+    if isinstance(ax, Axis):
+        return ax.size
+    else:
+        return prod(axis.size for axis in ensure_tuple(ax))  # type: ignore
+
+
 class dslice(eqx.Module):
     """
     Dynamic slice, comprising a (start, length) pair. Also aliased as ds.
@@ -524,6 +536,7 @@ __all__ = [
     "PartialShapeDict",
     "ShapeDict",
     "axis_name",
+    "axis_size",
     "concat_axes",
     "union_axes",
     "axis_spec_to_shape_dict",

{haliax-1.4.dev296 → haliax-1.4.dev298}/src/haliax/core.py

@@ -441,7 +441,11 @@ class NamedArray:
 
     @typing.overload
     def dot(
-        self, *args, axis: Optional[AxisSelection], precision: PrecisionLike = None, dot_general=jax.lax.dot_general
+        self,
+        *args: "NamedArray",
+        axis: Optional[AxisSelection],
+        precision: PrecisionLike = None,
+        dot_general=jax.lax.dot_general,
     ) -> "NamedArray":
         ...
 
@@ -1143,7 +1147,7 @@ def flatten_axes(array: NamedArray, old_axes: AxisSelection, new_axis: AxisSelec
     """
     old_axes = ensure_tuple(old_axes)
     old_axes = array.resolve_axis(old_axes)
-    total_axis_size = prod(array.axis_size(ax) for ax in old_axes)
+    total_axis_size = haliax.axis_size(old_axes)
 
     if isinstance(new_axis, Axis):
         if new_axis.size != total_axis_size:

{haliax-1.4.dev296 → haliax-1.4.dev298}/src/haliax/nn/embedding.py

@@ -1,4 +1,6 @@
 import dataclasses
+import math
+import warnings
 from typing import Optional
 
 import equinox as eqx
@@ -21,9 +23,14 @@ class Embedding(eqx.Module):
     Embed: AxisSpec = eqx.static_field()
 
     @staticmethod
-    def init(Vocab: Axis, Embed: AxisSpec, initializer_range: float = 0.02, *, key):
+    def init(Vocab: Axis, Embed: AxisSpec, *, init_scale: float = 1, key, initializer_range: Optional[float] = None):
+        if initializer_range is not None:
+            warnings.warn("initializer_range is deprecated. Use init_std instead.", DeprecationWarning)
+            init_scale = initializer_range
+
         all_axes = (Vocab,) + ensure_tuple(Embed)
-        weight = hax.random.normal(key, all_axes) * initializer_range
+        output_size = hax.axis_size(Embed)
+        weight = hax.random.truncated_normal(key, all_axes, -3, 3) * (init_scale / math.sqrt(output_size))
         return Embedding(weight=weight, Vocab=Vocab, Embed=Embed)
 
     def __call__(self, input_ids, *, key: Optional[PRNGKeyArray] = None):

{haliax-1.4.dev296 → haliax-1.4.dev298}/src/haliax/nn/linear.py

@@ -1,3 +1,4 @@
+import math
 from typing import Callable, Optional
 
 import equinox as eqx
@@ -25,7 +26,14 @@ class Linear(eqx.Module):
 
     @staticmethod
     def init(
-        In: AxisSpec, Out: AxisSpec, *, key, use_bias=True, out_first: bool = False, dot_general=None
+        In: AxisSpec,
+        Out: AxisSpec,
+        *,
+        key,
+        use_bias=True,
+        out_first: bool = False,
+        dot_general=None,
+        init_scale: float = 1.0,
     ) -> "Linear":
         """
 
@@ -36,9 +44,11 @@ class Linear(eqx.Module):
             use_bias: bool: Whether to use a bias term
             out_first: bool: Whether to put output axes first in the weight matrix. out_first is how PyTorch does it.
             dot_general: Callable: The dot_general function to use. Defaults to jax.lax.dot_general. For fp8 or int8
+            init_scale: float: The scale to use for initialization. We scale init by 1/sqrt(Input.size)*init_scale
         """
         joint_spec = hax.concat_axis_specs(Out, In) if out_first else hax.concat_axis_specs(In, Out)
-        weight = hax.random.normal(key, joint_spec) * 0.02
+        input_size = hax.axis_size(In)
+        weight = hax.random.truncated_normal(key, joint_spec, -3, 3) * (init_scale / math.sqrt(input_size))
         bias = hax.zeros(Out) if use_bias else None
 
         if dot_general is None:

{haliax-1.4.dev296 → haliax-1.4.dev298}/src/haliax/nn/mlp.py

@@ -48,6 +48,7 @@ class MLP(eqx.Module):
         use_final_bias: bool = True,
         key: PRNGKeyArray,
         dot_general: Optional[DotGeneralOp] = None,
+        init_scale: float = 1.0,
     ):
         Width = _get_width(width)
         Width2 = Width.alias(Width.name + "2")
@@ -58,18 +59,34 @@ class MLP(eqx.Module):
 
         if depth == 0:
             # special case: no hidden layers
-            layers.append(Linear.init(Input, Output, use_bias=use_final_bias, key=keys[0], dot_general=dot_general))
+            layers.append(
+                Linear.init(
+                    Input, Output, use_bias=use_final_bias, key=keys[0], dot_general=dot_general, init_scale=init_scale
+                )
+            )
         else:
             # first hidden layer
-            layers.append(Linear.init(Input, Width, use_bias=use_bias, key=keys[0], dot_general=dot_general))
+            layers.append(
+                Linear.init(
+                    Input, Width, use_bias=use_bias, key=keys[0], dot_general=dot_general, init_scale=init_scale
+                )
+            )
             # middle hidden layers
             cur = Width
             next = Width2
             for i in range(1, depth):
-                layers.append(Linear.init(cur, next, use_bias=use_bias, key=keys[i], dot_general=dot_general))
+                layers.append(
+                    Linear.init(
+                        cur, next, use_bias=use_bias, key=keys[i], dot_general=dot_general, init_scale=init_scale
+                    )
+                )
                 cur, next = next, cur
             # final hidden layer
-            layers.append(Linear.init(cur, Output, use_bias=use_final_bias, key=keys[-1], dot_general=dot_general))
+            layers.append(
+                Linear.init(
+                    cur, Output, use_bias=use_final_bias, key=keys[-1], dot_general=dot_general, init_scale=init_scale
+                )
+            )
 
         return MLP(
             layers=tuple(layers),

{haliax-1.4.dev296 → haliax-1.4.dev298}/tests/test_fp8.py

@@ -1,8 +1,10 @@
+import chex
 import equinox as eqx
 import jax.numpy as jnp
 import jax.random as jrandom
 import jax.tree_util
 import numpy as np
+from chex import assert_trees_all_close
 
 import haliax as hax
 from haliax._src.fp8 import compute_scale
@@ -19,18 +21,20 @@ from haliax.quantization import (
 def test_fp8_is_reasonable():
     In = hax.Axis("In", 8)
     Out = hax.Axis("Out", 8)
-    linear = Linear.init(In, Out, key=jrandom.PRNGKey(0))
+    linear = Linear.init(In, Out, key=jrandom.PRNGKey(0), init_scale=0.1)
 
-    fp8_linear = Linear.init(In, Out, key=jrandom.PRNGKey(0), dot_general=hax.quantization.Fp8DotGeneralOp.init())
+    fp8_linear = Linear.init(
+        In, Out, key=jrandom.PRNGKey(0), dot_general=hax.quantization.Fp8DotGeneralOp.init(), init_scale=0.1
+    )
 
-    input = hax.random.normal(jrandom.PRNGKey(0), In) * 10
+    input = hax.random.normal(jrandom.PRNGKey(3), In)
     output = linear(input)
     fp8_output = fp8_linear(input)
 
     assert output.shape == fp8_output.shape
     assert output.dtype == fp8_output.dtype
 
-    assert jnp.allclose(output.array, fp8_output.array, atol=1e-1, rtol=1e-1)
+    assert_trees_all_close(output.array, fp8_output.array, atol=1e-2, rtol=5e-2)
 
 
 # https://github.com/google/flax/blob/6f2b08e024c2fd2f8cec42a6c82408cb35412319/tests/linen/linen_test.py#L1222
@@ -130,9 +134,9 @@ def test_fp_loop():
 def test_layer_splicing():
     key, init_key, random_key = jrandom.split(jrandom.PRNGKey(seed=123), 3)
     Input = hax.Axis("Input", 16)
-    Hidden = hax.Axis("Hidden", 16)
+    Hidden = hax.Axis("Hidden", 64)
     Output = hax.Axis("Output", 32)
-    mlp = hax.nn.MLP.init(Input, Output, Hidden, 3, key=init_key)
+    mlp = hax.nn.MLP.init(Input, Output, Hidden, 3, key=init_key, init_scale=0.1)
 
     mlp_q = fp8_linear_layers(mlp, Fp8Config())
     for layer in mlp_q.layers:
@@ -141,7 +145,7 @@ def test_layer_splicing():
     input = hax.random.normal(jrandom.PRNGKey(0), Input) * 10  # 10 so we don't underflow
     output = mlp(input)
     output_q = mlp_q(input)
-    assert jnp.allclose(output.array, output_q.array, atol=1e-3, rtol=1e-3)
+    chex.assert_trees_all_close(output.array, output_q.array, atol=1e-3, rtol=1e-3)
     assert not jnp.allclose(output_q.array, 0)  # don't want them to all underflow
 
     mlp_q = fp8_linear_layers(mlp, Fp8Config(targets="layers.0"))

{haliax-1.4.dev296 → haliax-1.4.dev298}/tests/test_hof.py

@@ -87,6 +87,38 @@ def test_scan_doesnt_scan_scalars():
     assert jnp.all(jnp.equal(selected.array, named1.array * 4.0))
 
 
+def test_scan_doesnt_scan_init():
+    Height = Axis("Height", 10)
+    named1 = hax.random.uniform(PRNGKey(0), (Height,))
+
+    init = jnp.arange(Height.size, dtype=jnp.float32)
+
+    def scan_fun(acc, z, x):
+        out = acc + z * x, x * z
+        return out
+
+    total, selected = hax.scan(scan_fun, Height)(init, 4.0, named1)
+
+    assert jnp.all(jnp.isclose(total, init + jnp.sum(named1.array * 4.0)))
+
+    # double check with named array init
+    total, selected = hax.scan(scan_fun, Height)(hax.named(init, "Height"), 4.0, named1)
+
+    assert jnp.all(jnp.isclose(total.array, init + jnp.sum(named1.array * 4.0)))
+
+    # now do fold
+    def fold_fun(acc, z, x):
+        return acc + z * x
+
+    total = hax.fold(fold_fun, Height)(init, 4.0, named1)
+
+    assert jnp.all(jnp.isclose(total, init + jnp.sum(named1.array * 4.0)))
+
+    total = hax.fold(fold_fun, Height)(hax.named(init, "Height"), 4.0, named1)
+
+    assert jnp.all(jnp.isclose(total.array, init + jnp.sum(named1.array * 4.0)))
+
+
 def test_reduce():
     Height = Axis("Height", 10)
     Width = Axis("Width", 3)

{haliax-1.4.dev296 → haliax-1.4.dev298}/tests/test_partitioning.py

@@ -108,26 +108,21 @@ def test_pjit_class_init_with_args():
 def test_infer_resource_partition_gda_bug():
     devices = jax.devices()
     with Mesh(np.array(devices).reshape(-1, 1), (ResourceAxis.DATA, ResourceAxis.MODEL)):
-        jax.config.update("jax_parallel_functions_output_gda", True)
-        try:
 
-            def foo():
-                return hax.zeros((Dim1, Dim2, Dim3))
+        def foo():
+            return hax.zeros((Dim1, Dim2, Dim3))
 
-            pjit_foo = named_jit(foo, resource_map)
-            r = pjit_foo()
-            assert r.axes == (Dim1, Dim2, Dim3)
+        pjit_foo = named_jit(foo, resource_map)
+        r = pjit_foo()
+        assert r.axes == (Dim1, Dim2, Dim3)
 
-            def bar(x):
-                return x
-
-            # this won't work with GDAs
-            pjit_bar = named_jit(bar, resource_map)
-            r = pjit_bar(r)
-            assert r.axes == (Dim1, Dim2, Dim3)
+        def bar(x):
+            return x
 
-        finally:
-            jax.config.update("jax_parallel_functions_output_gda", False)
+        # this won't work with GDAs
+        pjit_bar = named_jit(bar, resource_map)
+        r = pjit_bar(r)
+        assert r.axes == (Dim1, Dim2, Dim3)
 
 
 @skip_if_not_enough_devices(4)

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

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