haliax 1.4.dev321__tar.gz → 1.4.dev324__tar.gz

{haliax-1.4.dev321 → haliax-1.4.dev324}/PKG-INFO

@@ -1,11 +1,11 @@
 Metadata-Version: 2.3
 Name: haliax
-Version: 1.4.dev321
+Version: 1.4.dev324
 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/
+Project-URL: Documentation, https://haliax.readthedocs.io/en/latest/
 Author-email: David Hall <dlwh@cs.stanford.edu>
-License-File: LICENSE
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Science/Research
 Classifier: License :: OSI Approved :: Apache Software License
@@ -16,6 +16,7 @@ Requires-Python: >=3.10
 Requires-Dist: equinox>=0.10.6
 Requires-Dist: jaxtyping>=0.2.20
 Requires-Dist: jmp>=0.0.4
+Requires-Dist: safetensors>=0.4.3
 Provides-Extra: dev
 Requires-Dist: chex>=0.1.86; extra == 'dev'
 Requires-Dist: mkdocs-include-markdown-plugin; extra == 'dev'

{haliax-1.4.dev321 → haliax-1.4.dev324}/docs/indexing.md

@@ -72,7 +72,7 @@ can use [haliax.where][] for most of the same functionality, though.
 
 Before we continue, a note on shapes in JAX. Most JAX code will be used inside `jit`, which means that the sizes of all
 arrays must be determined at compile time (i.e. when JAX interprets your functions abstractly). This is a hard
-requirement in XLA. It might worked around one day, but it's the way things are for now.
+requirement in XLA.
 
 A consequence of this restriction is that certain indexing patterns aren't allowed in `jit`-ed JAX code:
 
@@ -97,10 +97,31 @@ f(jnp.arange(10), 2)
 ```
 
 This is a not-uncommon pattern in situations where you want to process a large array in chunks. In Haliax, we provide
-two solutions: [haliax.slice][] as well as dynamic slices ([haliax.dslice][] a.k.a. [haliax.ds][]).
+two solutions: [haliax.slice][] and dynamic slices ([haliax.dslice][] a.k.a. [haliax.ds][]).
 
 ## Dynamic Slices
 
+[haliax.slice][] is a convenience function that wraps `jax.lax.dynamic_slice` and allows you to slice an array with a
+dynamic start and size. This is useful for situations where you need to slice an array in a way that can't be determined
+at compile time. For example, the above example can be written as follows:
+
+```python
+import jax
+
+import haliax as hax
+
+N = hax.Axis("N", 10)
+q = hax.arange(N)
+
+@hax.named_jit
+def f(x, slice_size: int):
+    num_blocks = N.size // slice_size
+    def body(i, m):
+        return i + hax.mean(hax.slice(x, {"N": i * slice_size}, {"N": slice_size}))
+    jax.lax.fori_loop(0, num_blocks, body, 0.0)
+```
+
+
 In light of the requirement that all array sizes be known at compile time, Haliax provides both a simple [haliax.slice][]
 function, as well as [haliax.dslice][], which can be used with `[]`. The simple slice function is just a wrapper
 around [jax.lax.dynamic_slice][]] and not worth discussing here.

{haliax-1.4.dev321 → haliax-1.4.dev324}/docs/matmul.md

@@ -62,7 +62,7 @@ y = hax.dot(x, w, c, axis=(), out_axes=(D, ..., H))  # shape is (D, W, C, H), eq
 [haliax.einsum][] is at its best when you want to express a more complex tensor contraction.
 It is similar to [numpy.einsum](https://numpy.org/doc/stable/reference/generated/numpy.einsum.html)
 or [einops.einsum](https://einops.rocks/api/einsum/) in terms of syntax and behavior,
-but extended to work with named axes, including added flexibility that named axes provide.
+but extended to work with named axes, including the added flexibility that named axes provide.
 Our "flavor" of `einsum` is most similar to `einops.einsum`'s flavor, in that
 it supports long names for axes (like `"batch h w, h w channel -> batch channel"`)
 rather than the compact notation of `numpy.einsum` (like `"bhwc,hwc->bc"`).

haliax-1.4.dev324/docs/state-dict.md

@@ -0,0 +1,160 @@
+# Serialization
+
+Haliax supports serialization of modules (including any [equinox.Module][]) to and from PyTorch-compatible
+state dicts using the [safetensors](https://github.com/huggingface/safetensors) library. For details on
+how state dicts work in PyTorch, see the [PyTorch documentation](https://pytorch.org/docs/stable/notes/serialization.html#saving-and-loading-torch-nn-modules).
+
+A state dict is a Python dictionary that maps string keys to tensors. It is used to store the parameters
+of a model (though typically not the model's structure or hyperparameters). The keys are typically the names of the
+model's parameters, arranged as `.`-separated paths. For example, a model with a `conv1` layer might have a
+state dict with keys like `conv1.weight` and `conv1.bias`. Sequences of modules (e.g., for lists of layers) are
+serialize with keys like `layer.0.weight`, `layer.1.weight`, etc.
+
+Haliax uses the [safetensors](https://github.com/huggingface/safetensors) library to serialize state dicts. This
+library is a safer, more portable format developed by Hugging Face. Serializing a native PyTorch state dict requires
+PyTorch itself, and we want to avoid that dependency. Also, PyTorch uses pickles, which are in general not
+safe to deserialize from untrusted sources.
+
+This does mean that you can't directly load a Haliax state dict into PyTorch, but safetensors is lightweight and
+easy to use. Hugging Face natively supports it in their libraries.
+
+## Saving a State Dict
+
+To serialize a module to a Pytorch-compatible state dict, use the [haliax.state_dict.to_torch_compatible_state_dict][]
+function. This function takes a module and returns a state dict. To save the state dict to a file, use the
+[haliax.state_dict.save_state_dict][] function, which writes the state dict to a file in safetensor format.
+`to_torch_compatible_state_dict` flattens [haliax.nn.Linear] module input and output axis specs to a format that
+is compatible with PyTorch Linear modules (though `out_first=True` is necessary to match PyTorch's Linear module).
+
+```python
+import haliax
+import jax.random as jrandom
+
+# Create a module
+Heads = haliax.Axis("Heads", 8)
+Dim = haliax.Axis("Dim", 16)
+Out = haliax.Axis("Out", 5)
+module = haliax.nn.Linear.init(In=(Heads, Dim), Out=Out, key=jrandom.PRNGKey(0))
+
+# Serialize the module to a state dict
+state_dict = haliax.state_dict.to_torch_compatible_state_dict(module)
+
+# Save the state dict to a file
+haliax.state_dict.save_state_dict(state_dict, 'state_dict.safetensors')
+```
+
+Note that the state dict is saved in the [safetensors](https://github.com/huggingface/safetensors) format, which
+is a safer, more portable format developed by Hugging Face. To load a model from a state dict in PyTorch, you
+can use safetensors directly.
+
+```python
+import torch
+from safetensors.torch import load_model
+
+model = torch.nn.Linear(10, 5)
+
+# Load the state dict from a file
+state_dict = load_model(model, 'state_dict.safetensors')
+```
+
+## Loading a State Dict
+
+Similarly, you can load a state dict from a file using the [haliax.state_dict.load_state_dict][] function. This
+function reads a state dict from a file in safetensors format and returns a dictionary. To load the state dict
+into a module, use the [haliax.state_dict.from_torch_compatible_state_dict][] function.
+
+```python
+import haliax as hax
+import jax.random as jrandom
+
+# Create a module
+Heads = hax.Axis("Heads", 8)
+Dim = hax.Axis("Dim", 16)
+Out = hax.Axis("Out", 5)
+module = hax.nn.Linear.init(In=(Heads, Dim), Out=Out, key=jrandom.PRNGKey(0))
+
+# Load the state dict from a file
+state_dict = hax.state_dict.load_state_dict('state_dict.safetensors')
+
+# this will unflatten the state dict and load it into the module
+module = hax.state_dict.from_torch_compatible_state_dict(module, state_dict)
+```
+
+The `from_torch_compatible_state_dict` function will unflatten the state dict and load it into the module. Note
+that the module must have the same structure as the module that was serialized to the state dict. If the module
+structure has changed, you may need to manually update the state dict keys to match the new structure.
+
+
+## Customizing Serialization
+
+### Changing the State Dict Key Names
+
+If for some reason you want to use different names in the serialized state dict (e.g. because you
+chose to use different names from a Hugging Face implementation), you can extend your class from  [haliax.state_dict.ModuleWithStateDictSerialization][]
+and use `_state_dict_key_map` to rename keys. For instance, the `Gpt2Transformer` class in Levanter has this method:
+
+```python
+from typing import Optional
+from haliax.state_dict import ModuleWithStateDictSerialization
+
+class Gpt2Transformer(ModuleWithStateDictSerialization):
+    ...
+
+    def _state_dict_key_map(self) -> dict[str, Optional[str]]:
+        return {"blocks": "h"}
+```
+
+This says that the field called `blocks` in this class should be (de)serialized as `h`,
+because the Hugging Face GPT-2 implementation uses `h`, which is not very clear.
+You can also "flatten" the submodules of a field by using `None`.
+
+### Custom Serialization Logic
+
+If your modules need fancier special logic, you'll need to extend your class from `ModuleWithStateDictSerialization` and
+override the default functions `to_state_dict()` and `from_state_dict()`. It takes in and returns a modified
+[haliax.state_dict.StateDict][]. As of June 2024, we almost never this in Levanter.
+
+For implementation, there are a few helper methods from `haliax.state_dict` that you can use:
+- To join specific prefix to the keys of Hugging Face state_dict, you can use the helper function `with_prefix()`.
+  The prefix comes from the name of attributes defined at the beginning of your model class.
+
+For example, below is the implementation of `to_state_dict()` in [levanter.models.backpack.BackpackLMHeadModel][].
+In this class, we want to preserve HF compatibility by saving untied output embeddings. (We chose not to implement
+non-weight-tied embeddings.)
+
+```python
+from typing import Optional
+
+from haliax.state_dict import with_prefix, StateDict
+
+
+class BackpackLMHeadModel(ModuleWithStateDictSerialization):
+    ...
+
+    def to_state_dict(self, prefix: Optional[str] = None) -> StateDict:
+        state_dict = super().to_state_dict(prefix=prefix)
+        # In levanter's implementation, we have a shared embedding matrix for both the word
+        # embeddings and the sense embeddings
+        state_dict[with_prefix(prefix, "backpack.word_embeddings.weight")] = state_dict[
+            with_prefix(prefix, "backpack.gpt2_model.wte.weight")
+        ]
+        state_dict[with_prefix(prefix, "backpack.position_embeddings.weight")] = state_dict[
+            with_prefix(prefix, "backpack.gpt2_model.wpe.weight")
+        ]
+        return state_dict
+```
+
+Similarly, to load weights from the state dict, you might need to implement `from_state_dict`. This function
+takes in a state dict and the module with the updated weights. You can use the `with_prefix()` helper function
+to join the prefix to the keys of the state dict.
+
+```python
+    def from_state_dict(self, state_dict: StateDict, prefix: Optional[str] = None) -> T:
+      ...
+
+
+```
+
+## API Reference
+
+::: haliax.state_dict

{haliax-1.4.dev321 → haliax-1.4.dev324}/mkdocs.yml

@@ -88,5 +88,6 @@ nav:
   - Partitioning: 'partitioning.md'
   - Higher Order Functions: 'hof.md'
   - FP8: 'fp8.md'
+  - Serialization: 'state-dict.md'
   - API Reference: 'api.md'
   - FAQ: 'faq.md'

{haliax-1.4.dev321 → haliax-1.4.dev324}/pyproject.toml

@@ -26,6 +26,7 @@ dependencies = [
     "equinox>=0.10.6",
     "jaxtyping>=0.2.20",
     "jmp>=0.0.4",
+    "safetensors>=0.4.3"
 ]
 dynamic =[ "version" ]
 
@@ -68,3 +69,4 @@ src_paths = ["src", "tests"]
 [project.urls]
 "Homepage" = "https://github.com/stanford-crfm/haliax"
 "Bug Tracker" = "https://github.com/stanford-crfm/haliax/issues/"
+"Documentation" = "https://haliax.readthedocs.io/en/latest/"

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

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

{haliax-1.4.dev321 → haliax-1.4.dev324}/src/haliax/__init__.py

@@ -14,6 +14,7 @@ import haliax.debug as debug
 import haliax.nn as nn
 import haliax.quantization as quantization
 import haliax.random as random
+import haliax.state_dict as state_dict
 import haliax.tree_util as tree_util
 import haliax.util as util
 
@@ -890,6 +891,7 @@ __all__ = [
     "random",
     "tree_util",
     "nn",
+    "state_dict",
     "Axis",
     "AxisSpec",
     "AxisSelection",

{haliax-1.4.dev321 → haliax-1.4.dev324}/src/haliax/_src/einsum.py

@@ -9,6 +9,7 @@ import haliax
 from ..axis import Axis, AxisSelector, axis_name, eliminate_axes, rearrange_for_partial_order, union_axes
 from ..core import NamedArray
 from ..jax_utils import _jittable_dg_einsum
+from ..quantization import DotGeneralOp
 from ..types import DTypeLike, PrecisionLike
 from ..util import ensure_tuple
 from .parsing import AliasTable, parse_einsum, raise_parse_error
@@ -19,8 +20,8 @@ def einsum(
     *arrays: NamedArray,
     precision: PrecisionLike = None,
     preferred_element_type: Optional[DTypeLike] = None,
-    _dot_general=jax.lax.dot_general,
-    **axis_aliases,
+    _dot_general: DotGeneralOp = jax.lax.dot_general,
+    **axis_aliases: AxisSelector,
 ) -> NamedArray:
     """Compute the tensor contraction of the input arrays according to Haliax's named variant of the Einstein summation
     convention.