search

stanford-crfm / haliax

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

Fix pspec method for levanter #88

Closed blahBlahhhJ closed 1 month ago

blahBlahhhJ commented 1 month ago

Handles the case when one logical axis maps to a tuple of (multiple) physical axes.