jax

b03ab674 - Replace uses of `.vma` (ShapedArray property) in favor of `.manual_axis_type.varying`.