No public description

PiperOrigin-RevId: 580799521
google-research · Nov 9, 2023 · efcf732 · efcf732
1 parent a9b35bc
commit efcf732
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/vmoe/nn/vit_moe_test.py b/vmoe/nn/vit_moe_test.py
@@ -86,7 +86,7 @@ def test_position_embedding_learned(self):
             'encoder_norm': {'bias': jnp.zeros(8), 'scale': jnp.ones(8)},
         },
     }, x)
-    expected_y = jax.nn.normalize(
+    expected_y = jax.nn.standardize(
         jnp.concatenate([pos_embedding, pos_embedding]), axis=-1)
     chex.assert_trees_all_close(y, expected_y)
 
@@ -124,7 +124,7 @@ def test_position_embedding_sincos2d(self, seq_len, pos_embedding):
         'params': {'encoder_norm': {'bias': jnp.zeros(8), 'scale': jnp.ones(8)}}
     }, x)
     pos_embedding = jnp.asarray(pos_embedding)
-    expected_y = jax.nn.normalize(
+    expected_y = jax.nn.standardize(
         jnp.concatenate([pos_embedding, pos_embedding]), axis=-1)
     chex.assert_trees_all_close(y, expected_y, rtol=1e-3)