search

stanford-crfm / haliax

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

Lowerable Named Jit #39

Closed dlwh closed 10 months ago

dlwh commented 10 months ago

Make named_jit support lower, which means we can ask how much memory it's going to need, etc.

Also bumps mypy and jax version for tests