jax

0a72e856 - Add **experimental** `with_dll_constraint` API. This is for cases when the users wants to let SPMD decide the sharding.

Add **experimental** `with_dll_constraint` API. This is for cases when the users wants to let SPMD decide the sharding. But this is a contradiction since layouts apply to device local shape and without knowing the sharding, you can't decide the layout. But there are cases where you don't care what the sharding is, you just want to force a row-major layout (for example). **This API should only be used for those cases**. PiperOrigin-RevId: 744888557