Fix incorrect matmul op in jax spectral normalization. Resolves #534

PiperOrigin-RevId: 470364802

Author: fehiepsi

edward2/jax/nn/normalization.py

@@ -140,7 +140,7 @@ def __call__(self, inputs: Array, training: bool = True) -> Array:
 
     if self.kernel_apply_kwargs is None:
       # By default, we use the implementation in SN-GAN.
-      kernel_apply = lambda x: w.reshape(-1, w.shape[-1]) @ x
+      kernel_apply = lambda x: x @ w.reshape(-1, w.shape[-1])
       in_shape = (np.prod(w.shape[:-1]),)
     else:
       # Otherwise, we extract the actual kernel transformation in the input

edward2/jax/nn/normalization_test.py

@@ -50,7 +50,7 @@ def setUp(self):
     self.norm_multiplier = 0.95
 
   @parameterized.named_parameters(
-      ("Dense", (None, 10), DenseLayer, ed.nn.SpectralNormalization),
+      ("Dense", (None, 3), DenseLayer, ed.nn.SpectralNormalization),
       ("Conv2D",
        (None, 32, 32, 3), Conv2DLayer, ed.nn.SpectralNormalizationConv2D))
   def test_spec_norm_magnitude(self, input_shape, layer, norm_wrapper):