haliax 1.4.dev394__tar.gz → 1.4.dev396__tar.gz

{haliax-1.4.dev394 → haliax-1.4.dev396}/PKG-INFO

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

@@ -2,16 +2,16 @@
 
 See also the [Equinox FAQ](https://docs.kidger.site/equinox/faq/)
 
-## Tip 1: `hax.debug.diagnose_common_issues`
+## [`hax.debug.diagnose_common_issues`][haliax.debug.diagnose_common_issues]
 
-`hax.debug.diagnose_common_issues` is a function that will raise an exception if it detects problems with your module.
+[haliax.debug.diagnose_common_issues][] is a function that will raise an exception if it detects problems with your module.
 Currently, we diagnose:
 
 * Reuse of arrays or NamedArrays in a field. [Equinox modules must be trees.](https://docs.kidger.site/equinox/faq/#a-module-saved-in-two-places-has-become-two-independent-copies)
 * Use of arrays or NamedArrays in a static field. Static data in JAX/Equinox must be hashable, and arrays are not hashable.
 
-## Tip 2: `hax.debug.visualize_shardings`
+## [`hax.debug.visualize_shardings`][haliax.debug.visualize_shardings]
 
-Use `hax.debug.visualize_shardings` to quickly inspect how a PyTree is sharded.
+Use [haliax.debug.visualize_shardings][] to quickly inspect how a PyTree is sharded.
 It prints the sharding of each array leaf, including the mapping from named axes
 to physical axes for :class:`haliax.NamedArray` leaves.

{haliax-1.4.dev394 → haliax-1.4.dev396}/docs/fp8.md

@@ -115,7 +115,7 @@ module = hax.quantization.apply_updates(module, updates, grads)
 ```
 
 That's it! Just a few lines of code to enable FP8. The `quantize_linear_layers` function will transform your module to use
-quantization-aware training for linear layers (or a subset if you want), and the combo of `partition_for_grad_overwrite` and `apply_updates` function will apply the updates to the module
+quantization-aware training for linear layers (or a subset if you want), and the combo of [haliax.quantization.partition_for_grad_overwrite][] and [haliax.quantization.apply_updates][] function will apply the updates to the module
 in a way that is compatible with FP8.
 
 ## How FP8 works
@@ -140,7 +140,7 @@ depend on the gradients.)
 The way this happens is by "hijacking" the gradient computation. When you call `eqx.filter_grad(loss_fn)(module, data)`,
 you will get the gradient computation as normal, but you'll also get the updated state of the FP8 `dot_general` module.
 This updated state needs to directly replace the state in the module (rather than be used for a gradient step), which is
-why you need to use the `partition_for_grad_overwrite`
+why you need to use the [haliax.quantization.partition_for_grad_overwrite][]
 
 The FP8 `dot_general` module is implemented in [haliax.quantization.Fp8DotGeneralOp][]. It's actually not that complicated:
 

{haliax-1.4.dev394 → haliax-1.4.dev396}/docs/indexing.md

@@ -101,7 +101,7 @@ two solutions: [haliax.slice][] and dynamic slices ([haliax.dslice][] a.k.a. [ha
 
 ## Dynamic Slices
 
-[haliax.slice][] is a convenience function that wraps `jax.lax.dynamic_slice` and allows you to slice an array with a
+[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:
 
@@ -124,7 +124,7 @@ def f(x, slice_size: int):
 
 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.
+around [jax.lax.dynamic_slice][] and not worth discussing here.
 
 `dslice` is a trick borrowed from the new experimental [jax.experimental.pallas][] module. It's essentially a slice,
 except that instead of a start and an end (and maybe a stride), it takes a start and a size. The size must be

{haliax-1.4.dev394 → haliax-1.4.dev396}/docs/matmul.md

@@ -11,7 +11,7 @@ more suitable for expressing a particular contraction In general:
 See also the API reference for [haliax.dot][] and [haliax.einsum][] and the
 [cheat sheet section](cheatsheet.md#matrix-multiplication).
 
-### `haliax.dot`
+### [`haliax.dot`][haliax.dot]
 
 With [haliax.dot][], you specify the axes to contract over, without needing to write out the
 axes you want to keep (though you can if you want):
@@ -57,7 +57,7 @@ y = hax.dot(x, w, c, axis=())  # shape is (H, W, D, C), equivalent to np.einsum(
 y = hax.dot(x, w, c, axis=(), out_axes=(D, ..., H))  # shape is (D, W, C, H), equivalent to np.einsum("hwdc,dc,c->dwch", x, w, c)
 ```
 
-### `haliax.einsum`
+### [`haliax.einsum`][haliax.einsum]
 
 [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)

{haliax-1.4.dev394 → haliax-1.4.dev396}/docs/partitioning.md

@@ -80,7 +80,7 @@ with hax.axis_mapping({"batch": "data"}):
 ```
 
 Unlike in JAX, which has separate APIs for partitioning arrays inside and outside of `jit`, Haliax has a single API:
-`hax.shard` work inside and outside of `jit`. Haliax automatically
+[haliax.shard][] works inside and outside of `jit`. Haliax automatically
 chooses which JAX function to use based on context.
 
 

{haliax-1.4.dev394 → haliax-1.4.dev396}/docs/scan.md

@@ -17,11 +17,11 @@ def scan(f, init, xs, length=None):
   return carry, np.stack(ys)
 ```
 
-Haliax provides two versions of this pattern: [haliax.fold][] and [haliax.scan][]. haliax.scan works much like JAX's scan,
-except it is curried and it works with NamedArrays. haliax.fold is a more restricted version of scan that is easier to
+Haliax provides two versions of this pattern: [haliax.fold][] and [haliax.scan][]. [haliax.scan][] works much like JAX's scan,
+except it is curried and it works with NamedArrays. [haliax.fold][] is a more restricted version of scan that is easier to
 use if you don't need the full generality of scan. (It works with functions that only return `carry`, not `carry, output`.)
 
-## `haliax.scan`
+## [haliax.scan][haliax.scan]
 
 Unlike JAX's scan, Haliax's scan is curried - it takes the function and configuration first, then the initial carry and scan arguments as a separate call: `scan(f, axis)(init, xs)`.
 
@@ -32,7 +32,7 @@ Unlike JAX's scan, Haliax's scan is curried - it takes the function and configur
 
 ### Basic Example
 
-Here's a practical example of using `haliax.scan` to sum values along an axis while keeping track of intermediates:
+Here's a practical example of using [haliax.scan][] to sum values along an axis while keeping track of intermediates:
 
 ```python
 Time = Axis("Time", 100)
@@ -89,9 +89,9 @@ final_state, path = hax.scan(simulate_brownian_motion, Time)(init_state, None)
 
 More commonly, you might use this for an RNN or Transformer model. (See [haliax.nn.Stacked][].)
 
-## `haliax.fold`
+## [haliax.fold][haliax.fold]
 
-`haliax.fold` is a simpler version of `haliax.scan` that is easier to use when you don't need the full generality of `scan`.
+[haliax.fold][] is a simpler version of [haliax.scan][] that is easier to use when you don't need the full generality of `scan`.
 Specifically, `fold` is for functions that only return a `carry`, not a `carry, output`.
 
 Morally, `fold` is like this Python code:
@@ -138,10 +138,10 @@ init_state = (
 final_state = hax.fold(running_stats, Time)(init_state, data)
 ```
 
-## `haliax.map`
+## [haliax.map][haliax.map]
 
-`haliax.map` is a convenience function that applies a function to each element of an axis. It is similar
-to [jax.lax.map][] but works with NamedArrays, providing a similar interface to `haliax.scan` and `haliax.fold`.
+[haliax.map][] is a convenience function that applies a function to each element of an axis. It is similar
+to [jax.lax.map][] but works with NamedArrays, providing a similar interface to [haliax.scan][] and [haliax.fold][].
 
 ```python
 
@@ -155,13 +155,13 @@ def my_fn(x):
 result = hax.map(my_fn, Time)(data)
 ```
 
-You should generally prefer to use [haliax.vmap][] instead of `haliax.map`, but it's there if you need it.
-(It uses less memory than `haliax.vmap` but is slower.)
+You should generally prefer to use [haliax.vmap][] instead of [haliax.map][], but it's there if you need it.
+(It uses less memory than [haliax.vmap][] but is slower.)
 
 
 ## Gradient Checkpointing / Rematerialization
 
-Both `haliax.scan` and `haliax.fold` support gradient checkpointing, which can be useful for deep models.
+Both [haliax.scan][] and [haliax.fold][] support gradient checkpointing, which can be useful for deep models.
 Typically, you'd use this as part of [haliax.nn.Stacked][] or [haliax.nn.BlockSeq][] but you can also use it directly.
 
 Gradient checkpointing is a technique for reducing memory usage during backpropagation by recomputing some
@@ -409,7 +409,7 @@ JAX array will have its first axis vmapped over.
 
 Sometimes you may want to apply each block independently, without feeding the
 output of one block into the next.  `Stacked.vmap` does exactly that: it uses
-[`haliax.vmap`][] to broadcast the initial value to every block and evaluates
+[haliax.vmap][] to broadcast the initial value to every block and evaluates
 them in parallel, returning the stack of outputs.
 
 ```python

haliax-1.4.dev396/docs/vmap.md

@@ -0,0 +1,38 @@
+## Vectorization with `haliax.vmap`
+
+`haliax.vmap` is a [`NamedArray`][haliax.NamedArray] aware wrapper around
+[`jax.vmap`][jax.vmap].  Instead of supplying positional axis numbers you pass
+the [`Axis`][haliax.Axis] (or axis name) you want to map over.  Any
+`NamedArray` containing that axis is mapped in parallel and the axis is
+reinserted in the output.  Regular JAX arrays can be mapped as well by
+providing a `default` spec or per‑argument overrides.
+
+Unlike vanilla `jax.vmap`, you may supply **one or more axes**.  When multiple
+axes are given, the function is vmapped over each axis in turn (innermost first).
+If an axis isn't already present in the array you must also specify its size,
+either by passing an `Axis` object (`Axis("batch", 4)`) or a mapping such as
+`{"batch": 4}` so the new dimension can be inserted.
+
+### Basic Example
+
+```python
+import haliax as hax
+
+Batch = hax.Axis("batch", 4)
+
+def double(x):
+    return x * 2
+
+x = hax.arange(Batch)
+y = hax.vmap(double, Batch)(x)
+```
+
+The result `y` has the same `Batch` axis as `x`, and each element was processed
+in parallel.  With JAX you would write `jax.vmap(double)(x.array)` and manually
+specify `in_axes`, but Haliax handles the axis automatically.
+
+For applying many modules in parallel see
+[`Stacked.vmap`](scan.md#apply-blocks-in-parallel-with-vmap) which builds on this
+primitive.
+
+::: haliax.vmap

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

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

haliax-1.4.dev394/docs/vmap.md

@@ -1,9 +0,0 @@
-## Vectorization
-
-
-This primitive is also used by [`Stacked.vmap`](scan.md#apply-blocks-in-parallel-with-vmap)
-to apply an entire stack of blocks in parallel.
-
-(This is a work in progress. Please contact dlwh for more information.)
-
-::: haliax.vmap

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

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