Issues
- 8
- 17
Move `state_dict` serialization to Haliax
#35 opened by dlwh - 3
Cleaner All Pairs Difference
#106 opened by 0xc1c4da - 2
- 5
make Linear support overlapping input/axis names?
#53 opened by dlwh - 2
AxisSpec should allow nested sequences--and all functions that consume an AxisSpec should support nested sequences.
#107 opened by cooljoseph1 - 2
Proposal: Make haliax.nn modules have an __init__ function which can be used to construct a skeleton tree
#104 opened by cooljoseph1 - 0
Add a tutorial about Stacked/scan layers
#103 opened by dlwh - 1
Make a `scan_aware_tree_map`
#99 opened by dlwh - 1
Use Generics
#97 opened by thomasahle - 5
- 5
- 2
jax.random.KeyArray deprecated
#75 opened by jennifgcrl - 0
- 10
- 0
- 1
Make a helper function to squash/unsquash all axes (except some) into a single batch axis
#67 opened by dlwh - 1
Lp norm function?
#64 opened by rohan-mehta-1024 - 0
Add more NN Layers
#22 opened by dlwh - 0
auto_shard inside zeros etc
#62 opened by dlwh - 0
better stack trace for concretization inside scan
#57 opened by dlwh - 0
better exception when calling .item inside jit
#58 opened by dlwh - 1
- 2
Support Rearrange Without Tailing Dimensions
#52 opened by karan-dalal - 5
Model Inference with Haliax?
#10 opened by rohan-mehta-1024 - 1
NN Layers: Add Pooling
#25 opened by dlwh - 0
Adaptive Pooling
#50 opened by dlwh - 2
`jax.lax.switch` for Haliax
#46 opened by anh-tong - 1
Add einops-style `rearrange`
#19 opened by dlwh - 3
How does haliax work with mixed precision?
#43 opened by samuelstevens - 3
Worth adding top_k function?
#32 opened by rohan-mehta-1024 - 0
NN Layers: Convolutions
#23 opened by dlwh - 0
RNN: Add RNN Cells
#27 opened by dlwh - 1
In-place Array Modificatons
#36 opened by bmac3 - 3
Add `ravel` / `flatten`
#18 opened by dlwh - 0
- 0
- 0
NN Layers: RNN Backbone
#26 opened by dlwh - 0
NN Layers: Add GroupNorm
#24 opened by dlwh - 0
add hypothesis tests (mostly testing shapes)
#21 opened by dlwh - 4
Make Haliax docs
#4 opened by dlwh - 2
Why init instead of __init__?
#17 opened by rohan-mehta-1024 - 3
Why do all axes have to be unique?
#13 opened by rohan-mehta-1024 - 2
Tutorial Suggestion
#8 opened by pharringtonp19 - 0
Equinox woes: declared class attributes
#3 opened by dlwh - 1
Pipeline Parallelism in `Stacked`
#1 opened by dlwh