haliax 1.4.dev395__tar.gz → 1.4.dev396__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev395
+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.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.dev395/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.dev395/src/haliax/__about__.py

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