[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: 480841853
Change-Id: I18a744a561d73da3ccd43c26fcd99d10070c058e

alphafold/model/layer_stack_test.py

@@ -22,7 +22,7 @@ import haiku as hk
 import jax
 import jax.numpy as jnp
 import numpy as np
-import scipy
+import scipy.stats
 
 
 # Suffixes applied by Haiku for repeated module names.