stanford-crfm / haliax

Named Tensors for Legible Deep Learning in JAX
Apache License 2.0
140 stars 9 forks source link

add in-place mutation ops for NamedArrays, support lists/1-d JAX Arrays for indexing #85

Closed dlwh closed 2 months ago