Minor fix in SpectralNormalization to ensure sigma is scalar.

PiperOrigin-RevId: 351168514

Author: jtainslie

edward2/tensorflow/layers/normalization.py

@@ -377,6 +377,8 @@ def update_weights(self):
         u_hat = tf.nn.l2_normalize(tf.matmul(v_hat, w_reshaped))
 
     sigma = tf.matmul(tf.matmul(v_hat, w_reshaped), tf.transpose(u_hat))
+    # Convert sigma from a 1x1 matrix to a scalar.
+    sigma = tf.reshape(sigma, [])
     u_update_op = self.u.assign(u_hat)
     v_update_op = self.v.assign(v_hat)