Satisfy stricter input requirement for jnp.linalg.norm

The previous version of the code caused an error under the new more restrictive API of `jnp.linalg.norm` (see jax-ml/jax#12670).

PiperOrigin-RevId: 484080138

jax_cfd/base/initial_conditions.py

@@ -65,7 +65,7 @@ def _log_normal_pdf(x, mode, variance=.25):
 
 
 def _max_speed(v):
-  return jnp.linalg.norm([u.data for u in v], axis=0).max()
+  return jnp.linalg.norm(jnp.array([u.data for u in v]), axis=0).max()
 
 
 def filtered_velocity_field(

jax_cfd/base/initial_conditions_test.py

@@ -50,7 +50,7 @@ def test_filtered_velocity_field(
       self, seed, grid, maximum_velocity, peak_wavenumber):
     v = ic.filtered_velocity_field(
         jax.random.PRNGKey(seed), grid, maximum_velocity, peak_wavenumber)
-    actual_maximum_velocity = jnp.linalg.norm([u.data for u in v], axis=0).max()
+    actual_maximum_velocity = jnp.linalg.norm(jnp.array([u.data for u in v]), axis=0).max()
     max_divergence = fd.divergence(v).data.max()
 
     # Assert that initial velocity is divergence free

jax_cfd/collocated/initial_conditions.py

@@ -44,7 +44,7 @@ def _log_normal_pdf(x, mode, variance=.25):
 
 
 def _max_speed(v):
-  return jnp.linalg.norm([u.data for u in v], axis=0).max()
+  return jnp.linalg.norm(jnp.array([u.data for u in v]), axis=0).max()
 
 
 def filtered_velocity_field(

jax_cfd/collocated/initial_conditions_test.py

@@ -43,7 +43,7 @@ def test_filtered_velocity_field(
       self, seed, grid, maximum_velocity, peak_wavenumber):
     v = initial_conditions.filtered_velocity_field(
         jax.random.PRNGKey(seed), grid, maximum_velocity, peak_wavenumber)
-    actual_maximum_velocity = jnp.linalg.norm([u.data for u in v], axis=0).max()
+    actual_maximum_velocity = jnp.linalg.norm(jnp.array([u.data for u in v]), axis=0).max()
     max_divergence = fd.centered_divergence(v).data.max()
 
     # Assert that initial velocity is divergence free