jax

b1843782 - Add a private parameter `_manual_axis_type` on `lax.empty` so that pl.kernel/core_map can propagate `manual_axis_type` on refs correctly.

Add a private parameter `_manual_axis_type` on `lax.empty` so that pl.kernel/core_map can propagate `manual_axis_type` on refs correctly. Maybe `empty_ref` would also need `out_sharding` and `manual_axis_type` parameters on it but we can leave that as a TODO. PiperOrigin-RevId: 902756246