jax

25094923 - Add a batching rule for with_layout_constraint by making the new dimension that vmap inserts major-most.