[JAX] Import scipy submodules directly in code that uses them.
[JAX] Import scipy submodules directly in code that uses them.

JAX will soon import scipy lazily, so code that relies on JAX having already imported submodules of scipy will fail. Change code to import what it uses.

PiperOrigin-RevId: 479652259

google/jax-cfd