stanford-crfm haliax issues

stanford-crfm / haliax

Named Tensors for Legible Deep Learning in JAX

Apache License 2.0

140 stars 9 forks source link

issues

Newest

Newest Most commented Recently updated Oldest Least commented Least recently updated

Use Generics

#97 thomasahle opened 2 days ago
1
Another attempt at shard

#96 dlwh closed 1 week ago
0
Run levanter tests

#95 dlwh closed 3 weeks ago
0
No default in_resources

#94 dlwh closed 3 weeks ago
1
Use sqrt(fan_in) as default in and truncated_normal

#93 dlwh closed 3 weeks ago
0
scan/fold should not consider the `init` args for scanning over

#92 rjpower closed 4 weeks ago
5
Expose the filter_checkpoint helper: it's useful for custom training loops.

#91 rjpower closed 1 month ago
0
Move StateDict to Haliax, clean it up a lot.

#90 dlwh opened 1 month ago
0
revert previous PR

#89 blahBlahhhJ closed 1 month ago
0
Fix pspec method for levanter

#88 blahBlahhhJ closed 1 month ago
0
Add replica axes

#87 blahBlahhhJ closed 1 month ago
1
fix hax.where in the presence of scalar NamedArrays

#86 dlwh closed 2 months ago
0
add in-place mutation ops for NamedArrays, support lists/1-d JAX Arrays for indexing

#85 dlwh closed 2 months ago
0
allow einsum to accept aliases similar to rearrange

#84 dlwh closed 2 months ago
0
switch to pdm:

#83 dlwh closed 2 months ago
0
fix revision number calculation for auto-publishing

#82 dlwh closed 2 months ago
0
Poetry + Dev Builds

#81 dlwh closed 3 months ago
0
Fix multi device shard

#80 dlwh closed 3 months ago
0
Simplify sharding

#79 dlwh closed 3 months ago
0
Make Conv work when input type is not the same as kernel type

#78 dlwh opened 3 months ago
0
Add FP8 support to Haliax

#77 dlwh closed 3 months ago
0
No Scan Layers

#76 dlwh closed 3 months ago
0
jax.random.KeyArray deprecated

#75 jennifgcrl closed 2 months ago
2
workaround current limitation in jax_metals

#74 dlwh closed 4 months ago
0
explicitly shard scalars in named_jit

#73 dlwh closed 4 months ago
0
Support for partitioning/sharded data with Pallas kernels?

#72 G-Levine opened 4 months ago
10
tile and repeat

#71 dlwh closed 4 months ago
0
add a non-scan-layers version of Stacked

#70 dlwh closed 3 months ago
0
Add single argument mode for hax.where

#69 blahBlahhhJ closed 4 months ago
0
Don't check shapes in tree_unflatten. Fixes #66

#68 dlwh closed 4 months ago
0
Make a helper function to squash/unsquash all axes (except some) into a single batch axis

#67 dlwh closed 4 months ago
1
named arrays and `eqxi.while_loop(mode=checkpointed)` do not get a long

#66 dlwh closed 4 months ago
0
Mixed Precision and Resource Envs

#65 dlwh opened 5 months ago
0
Lp norm function?

#64 rohan-mehta-1024 opened 5 months ago
1
Einsum

#63 dlwh closed 5 months ago
0
auto_shard inside zeros etc

#62 dlwh opened 5 months ago
0
top_k returns indices in addition to values

#61 blahBlahhhJ closed 5 months ago
1
Add checkpointed_scan

#60 dlwh opened 5 months ago
0
add htu.tree_map etc

#59 dlwh closed 5 months ago
0
better exception when calling .item inside jit

#58 dlwh opened 5 months ago
0
better stack trace for concretization inside scan

#57 dlwh opened 5 months ago
0
Dev

#55 dlwh closed 4 months ago
0
Add tree_map etc variants that set is_leaf=is_named_array by default

#54 dlwh closed 5 months ago
1
make Linear support overlapping input/axis names?

#53 dlwh opened 6 months ago
0
Support Rearrange Without Tailing Dimensions

#52 karan-dalal closed 5 months ago
2
Add Pooling operations

#51 dlwh closed 6 months ago
0
Adaptive Pooling

#50 dlwh opened 8 months ago
0
refactor dot_product_attention to be more general but also check argu…

#49 dlwh closed 8 months ago
0
Add a simple MLP module

#48 dlwh closed 8 months ago
0
Update README.md

#47 jennifgcrl closed 8 months ago
0