Package: haliax
HaliaxThough 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 WindHaliax is aJAXlibrary for building neural networks with named tensors, in the tradition of Alexander Rush'sTensor Considered Harmful.
Named tensors improve thelegibilityandcompositionalityof tensor programs by using named axes instead of positional indices
as typically used in NumPy, PyTorch, etc.Despite the focus on legibility, Haliax
is alsofast, typically about as fast as "pure" JAX code.
Haliax is also built to bescalable: it
can supportFully-Sharded Data Parallelism (FSDP)and Tensor Parallelism withjust a few lines of code. Haliax powersLevanter,
our companion library for training large language models and other foundation models, with scale proven up to 20B parameters
and up to a TPU v3-256 pod slice.Example: AttentionHere's a minimal attention module implementation in Haliax. For a more detailed introduction,
please see theHaliax tutorial.
(We use the excellentEquinoxlibrary for its module system and tree transformations.)importequinoxaseqximportjaximportjax.numpyasjnpimporthaliaxashaximporthaliax.nnashnnPos=hax.Axis("position",1024)# sequence lengthKPos=Pos.alias("key_position")Head=hax.Axis("head",8)# number of attention headsKey=hax.Axis("key",64)# key sizeEmbed=hax.Axis("embed",512)# embedding sizedefattention_scores(Key,KPos,query,key,mask):# how similar is each query to each keyscores=hax.dot(Key,query,key)/jnp.sqrt(Key.size)ifmaskisnotNone:scores-=1E9*(1.0-mask)# convert to probabilitiesscores=hax.nn.softmax(scores,KPos)returnscoresdefattention(Key,KPos,query,key,value,mask):scores=attention_scores(Key,KPos,query,key,mask)answers=hax.dot(KPos,scores,value)returnanswers# Causal Mask means that if pos >= key_pos, then pos can attend to key_poscausal_mask=hax.arange(Pos).broadcast_axis(KPos)>=hax.arange(KPos)classAttention(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]@staticmethoddefinit(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)returnAttention(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 attentionk=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)returnxHaliax was created byStanford's Center for Research on Foundation Models (CRFM)'s research engineering team.
You can find us in the #levanter channel on the unofficialJax LLM Discord.DocumentationTutorialsThese are some tutorials to get you started with Haliax. They are available as Colab notebooks:Introduction to Haliax with TransformersDistributed Training in Haliax(including FSDP)Tensor Parallelism in HaliaxMixed Precision withjmp(This one is really a tutorial forjmpbut it's how to use it with Haliax...)API ReferenceHaliax's API documentation is available athaliax.readthedocs.io.ContributingWe welcome contributions! Please seeCONTRIBUTING.mdfor more information.
We also have a list ofgood first issuesto help you get started. (If those don't appeal, don't hesitate to reach out to us on Discord!)LicenseHaliax is licensed under the Apache License, Version 2.0. SeeLICENSEfor the full license text.