This release is a pre-release and may not be stable for production use.
Haliax
Though you don’t seem to be much for listening, it’s best to be careful. If you managed to catch hold of even just a piece of my name, you’d have all manner of power over me.
— Patrick Rothfuss, The Name of the Wind
Haliax is a JAX library for building neural networks with named tensors, in the tradition of Alexander Rush's Tensor Considered Harmful. Named tensors improve the legibility and compositionality of tensor programs by using named axes instead of positional indices as typically used in NumPy, PyTorch, etc.
Despite the focus on legibility, Haliax is also fast, typically about as fast as "pure" JAX code. Haliax is also built to be scalable: it can support Fully-Sharded Data Parallelism (FSDP) and Tensor Parallelism with just a few lines of code. Haliax powers Levanter, our companion library for training large language models and other foundation models, with scale proven up to 70B parameters and up to TPU v4-2048.
Example: Attention
Here's a minimal attention module implementation in Haliax. For a more detailed introduction, please see the Haliax tutorial. (We use the excellent Equinox library for its module system and tree transformations.)
import equinox as eqx
import jax
import jax.numpy as jnp
import haliax as hax
import haliax.nn as hnn
Pos = hax.Axis("position", 1024) # sequence length
KPos = Pos.alias("key_position")
Head = hax.Axis("head", 8) # number of attention heads
Key = hax.Axis("key", 64) # key size
Embed = hax.Axis("embed", 512) # embedding size
# alternatively:
#Pos, KPos, Head, Key, Embed = hax.make_axes(pos=1024, key_pos=1024, head=8, key=64, embed=512)
def attention_scores(Key, KPos, query, key, mask):
# how similar is each query to each key
scores = hax.dot(query, key, axis=Key) / jnp.sqrt(Key.size)
if mask is not None:
scores -= 1E9 * (1.0 - mask)
# convert to probabilities
scores = haliax.nn.softmax(scores, KPos)
return scores
def attention(Key, KPos, query, key, value, mask):
scores = attention_scores(Key, KPos, query, key, mask)
answers = hax.dot(scores, value, axis=KPos)
return answers
# Causal Mask means that if pos >= key_pos, then pos can attend to key_pos
causal_mask = hax.arange(Pos).broadcast_axis(KPos) >= hax.arange(KPos)
class Attention(eqx.Module):
proj_q: hnn.Linear # [Embed] -> [Head, Key]
proj_k: hnn.Linear # [Embed] -> [Head, Key]
proj_v: hnn.Linear # [Embed] -> [Head, Key]
proj_answer: hnn.Linear # output projection from [Head, Key] -> [Embed]
@staticmethod
def init(Embed, Head, Key, *, key):
k_q, k_k, k_v, k_ans = jax.random.split(key, 4)
proj_q = hnn.Linear.init(In=Embed, Out=(Head, Key), key=k_q)
proj_k = hnn.Linear.init(In=Embed, Out=(Head, Key), key=k_k)
proj_v = hnn.Linear.init(In=Embed, Out=(Head, Key), key=k_v)
proj_answer = hnn.Linear.init(In=(Head, Key), Out=Embed, key=k_ans)
return Attention(proj_q, proj_k, proj_v, proj_answer)
def __call__(self, x, mask=None):
q = self.proj_q(x)
# Rename "position" to "key_position" for self attention
k = self.proj_k(x).rename({"position": "key_position"})
v = self.proj_v(x).rename({"position": "key_position"})
answers = attention(Key, KPos, q, k, v, causal_mask)
x = self.proj_answer(answers)
return x
Haliax was created by Stanford's Center for Research on Foundation Models (CRFM)'s research engineering team. You can find us in the #levanter channel on the unofficial Jax LLM Discord.
Documentation
Tutorials
These are some tutorials to get you started with Haliax. They are available as Colab notebooks:
- Introduction to Haliax with Transformers
- Distributed Training in Haliax (including FSDP)
- Tensor Parallelism in Haliax
- Mixed Precision with
jmp(This one is really a tutorial for jmp but it's how to use it with Haliax...)
API Reference
Haliax's API documentation is available at haliax.readthedocs.io.
Contributing
We welcome contributions! Please see CONTRIBUTING.md for more information. We also have a list of good first issues to help you get started. (If those don't appeal, don't hesitate to reach out to us on Discord!)
License
Haliax is licensed under the Apache License, Version 2.0. See LICENSE for the full license text.
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file marin_haliax-0.2.76.dev31155643335.tar.gz.
File metadata
- Download URL: marin_haliax-0.2.76.dev31155643335.tar.gz
- Upload date:
- Size: 779.4 kB
- Tags: Source
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/6.1.0 CPython/3.13.13
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
460aafccf4aab7a286841b77a4543c74b5b9aff42ea572c7a09995960c583fbd
|
|
| MD5 |
093ce5da972316fd1abf3baf77f7a2dc
|
|
| BLAKE2b-256 |
bd2888e466678d8ff083b732e501162197ad958f4ae860678184f208b24cb5a6
|
Provenance
The following attestation bundles were made for marin_haliax-0.2.76.dev31155643335.tar.gz:
Publisher:
marin-release-libs-wheels.yaml on marin-community/marin
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
marin_haliax-0.2.76.dev31155643335.tar.gz -
Subject digest:
460aafccf4aab7a286841b77a4543c74b5b9aff42ea572c7a09995960c583fbd - Sigstore transparency entry: 2366232053
- Sigstore integration time:
-
Permalink:
marin-community/marin@839e5e9e18299d750e6009b58f4535f77f9edafa -
Branch / Tag:
refs/heads/main - Owner: https://github.com/marin-community
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
marin-release-libs-wheels.yaml@839e5e9e18299d750e6009b58f4535f77f9edafa -
Trigger Event:
schedule
-
Statement type:
File details
Details for the file marin_haliax-0.2.76.dev31155643335-py3-none-any.whl.
File metadata
- Download URL: marin_haliax-0.2.76.dev31155643335-py3-none-any.whl
- Upload date:
- Size: 157.7 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/6.1.0 CPython/3.13.13
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
c4498229f88f484171f25cb8ff6883e8fcca7e600fa7115f6f4b3d19b3608267
|
|
| MD5 |
637050b2b6a294680a687c5a39ad6891
|
|
| BLAKE2b-256 |
5cc586797d00ca7e0b61208a8f0d04e2f2d95cea4abbdbf3f6ba89669d059b39
|
Provenance
The following attestation bundles were made for marin_haliax-0.2.76.dev31155643335-py3-none-any.whl:
Publisher:
marin-release-libs-wheels.yaml on marin-community/marin
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
marin_haliax-0.2.76.dev31155643335-py3-none-any.whl -
Subject digest:
c4498229f88f484171f25cb8ff6883e8fcca7e600fa7115f6f4b3d19b3608267 - Sigstore transparency entry: 2366234068
- Sigstore integration time:
-
Permalink:
marin-community/marin@839e5e9e18299d750e6009b58f4535f77f9edafa -
Branch / Tag:
refs/heads/main - Owner: https://github.com/marin-community
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
marin-release-libs-wheels.yaml@839e5e9e18299d750e6009b58f4535f77f9edafa -
Trigger Event:
schedule
-
Statement type:
