Use lax.slice_in_dim instead of lax.dynamic_slice_in_dim
Use lax.slice_in_dim instead of lax.dynamic_slice_in_dim

This should result in more efficient gradient calculations:
https://github.com/google/jax/issues/6872

PiperOrigin-RevId: 407151652

google/jax-cfd