[JAX] Make haiku tests pass with jax_explicit_x64_dtypes=ERROR

PiperOrigin-RevId: 914349194

Author: hawkinsp

haiku/_src/basic.py

@@ -166,6 +166,7 @@ def __call__(
     if not inputs.shape:
       raise ValueError("Input must not be scalar.")
 
+    inputs = jnp.asarray(inputs)
     input_size = self.input_size = inputs.shape[-1]
     output_size = self.output_size
     dtype = inputs.dtype

haiku/_src/conv.py

@@ -177,6 +177,7 @@ def __call__(
         unbatched, or an array of shape ``[N, spatial_dims, output_channels]``
         and rank-N+2 if batched.
     """
+    inputs = jnp.asarray(inputs)
     unbatched_rank = self.num_spatial_dims + 1
     allowed_ranks = [unbatched_rank, unbatched_rank + 1]
     if inputs.ndim not in allowed_ranks:
@@ -574,6 +575,7 @@ def __call__(
         unbatched, or an array of shape ``[N, spatial_dims, output_channels]``
         and rank-N+2 if batched.
     """
+    inputs = jnp.asarray(inputs)
     unbatched_rank = self.num_spatial_dims + 1
     allowed_ranks = [unbatched_rank, unbatched_rank + 1]
     if inputs.ndim not in allowed_ranks:

haiku/_src/layer_norm.py

@@ -149,6 +149,7 @@ def __call__(
     Returns:
       The array, normalized.
     """
+    inputs = jnp.asarray(inputs)
     if self.create_scale and  scale is not None:
       raise ValueError(
           "Cannot pass `scale` at call time if `create_scale=True`.")

haiku/_src/layer_norm_test.py

@@ -254,7 +254,7 @@ def test_multiple_param_axis(self, param_axis, param_shape):
     ln(x)
     self.assertEqual(ln.params_dict()["layer_norm/scale"].shape, param_shape)
     self.assertEqual(ln.params_dict()["layer_norm/offset"].shape, param_shape)
-  
+
   @staticmethod
   def _ln_f(axis, param_axis, create_scale=True, create_offset=True):
     def f(x):
@@ -269,7 +269,7 @@ def f(x):
   def test_layernorm_dtype_propagation_float64(self):
     """Input dtype -> output dtype; params dtype match input."""
     fwd = LayerNormTest._ln_f(axis=-1, param_axis=-1)
-    x = jnp.arange(12, dtype=jnp.float64).reshape(3, 4)
+    x = jnp.arange(12, dtype=float).reshape(3, 4)
     params = fwd.init(jax.random.PRNGKey(0), x)
     y = fwd.apply(params, None, x)
     self.assertEqual(y.dtype, x.dtype)

haiku/_src/nets/vqvae_test.py

@@ -113,17 +113,11 @@ def testNoneBatch(self, constructor, kwargs):
     vqvae_module(inputs, is_training=False)
 
   @parameterized.parameters({'use_jit': True, 'dtype': jnp.float32},
-                            {'use_jit': True, 'dtype': jnp.float64},
-                            {'use_jit': False, 'dtype': jnp.float32},
-                            {'use_jit': False, 'dtype': jnp.float64})
+                            {'use_jit': False, 'dtype': jnp.float32})
   @test_utils.transform_and_run
   def testEmaUpdating(self, use_jit, dtype):
-    if jax.local_devices()[0].platform == 'tpu' and dtype == jnp.float64:
-      self.skipTest('F64 not supported by TPU')
-
     embedding_dim = 6
-    np_dtype = np.float64 if dtype is jnp.float64 else np.float32
-    decay = np.array(0.1, dtype=np_dtype)
+    decay = np.array(0.1, dtype=np.float32)
     vqvae_module = vqvae.VectorQuantizerEMA(
         embedding_dim=embedding_dim,
         num_embeddings=7,

haiku/_src/pool.py

@@ -94,6 +94,7 @@ def max_pool(
   if padding not in ("SAME", "VALID"):
     raise ValueError(f"Invalid padding '{padding}', must be 'SAME' or 'VALID'.")
 
+  value = jnp.asarray(value)
   _warn_if_unsafe(window_shape, strides)
   window_shape = _infer_shape(value, window_shape, channel_axis)
   strides = _infer_shape(value, strides, channel_axis)
@@ -127,6 +128,7 @@ def avg_pool(
   if padding not in ("SAME", "VALID"):
     raise ValueError(f"Invalid padding '{padding}', must be 'SAME' or 'VALID'.")
 
+  value = jnp.asarray(value)
   _warn_if_unsafe(window_shape, strides)
   window_shape = _infer_shape(value, window_shape, channel_axis)
   strides = _infer_shape(value, strides, channel_axis)

haiku/_src/recurrent.py

@@ -260,6 +260,8 @@ def __init__(
     self.double_bias = double_bias
 
   def __call__(self, inputs, prev_state):
+    inputs = jnp.asarray(inputs)
+    prev_state = jnp.asarray(prev_state)
     input_to_hidden = hk.Linear(self.hidden_size)
     # TODO(b/173771088): Consider changing default to double_bias=False.
     hidden_to_hidden = hk.Linear(self.hidden_size, with_bias=self.double_bias)
@@ -329,6 +331,8 @@ def __call__(
       inputs: jax.Array,
       prev_state: LSTMState,
   ) -> tuple[jax.Array, LSTMState]:
+    inputs = jnp.asarray(inputs)
+    prev_state = jax.tree.map(jnp.asarray, prev_state)
     if len(inputs.shape) > 2 or not inputs.shape:
       raise ValueError("LSTM input must be rank-1 or rank-2.")
     x_and_h = jnp.concatenate([inputs, prev_state.hidden], axis=-1)
@@ -410,6 +414,8 @@ def __call__(
       inputs,
       state: LSTMState,
   ) -> tuple[jax.Array, LSTMState]:
+    inputs = jnp.asarray(inputs)
+    state = jax.tree.map(jnp.asarray, state)
     input_to_hidden = hk.ConvND(
         num_spatial_dims=self.num_spatial_dims,
         output_channels=4 * self.output_channels,
@@ -559,6 +565,8 @@ def __init__(
     self.b_init = b_init or jnp.zeros
 
   def __call__(self, inputs, state):
+    inputs = jnp.asarray(inputs)
+    state = jnp.asarray(state)
     if inputs.ndim not in (1, 2):
       raise ValueError("GRU input must be rank-1 or rank-2.")
 
@@ -650,6 +658,9 @@ def __call__(self, inputs, state):
       Tuple of the wrapped core's ``output, next_state``.
     """
     inputs, should_reset = inputs
+    inputs = jax.tree.map(jnp.asarray, inputs)
+    should_reset = jax.tree.map(jnp.asarray, should_reset)
+    state = jax.tree.map(jnp.asarray, state)
     if jax.tree_util.treedef_is_leaf(jax.tree.structure(should_reset)):
       # Equivalent to not tree.is_nested, but with support for Jax extensible
       # pytrees.

haiku/_src/rms_norm.py

@@ -106,6 +106,7 @@ def __call__(self, inputs: jax.Array):
     Returns:
       The normalized array, of the same shape as the inputs.
     """
+    inputs = jnp.asarray(inputs)
     axis = self.axis
     if isinstance(axis, slice):
       axis = tuple(range(inputs.ndim)[axis])