jax

e11e9d2e - Remove _manual_axis_type from `empty` now that it's unused given pl.kernel moved to mpmd_map