flâneur — a map of the web's best reading

marin-community/haliax: Named Tensors for Legible Deep Learning in JAX ·

github.com · 867 words · saved by 1 readers

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. Here's a minimal attention module implementation in Haliax. For

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

Explore this link on the map →

related reading