jax

989281dc - Add `to_ct_mt` on ManualAxisType to mirror PartitionSpec. Might be helpful for HiJAX (TBD though)