Merge pull request jax-ml#1323 from google/ortho-typo

fix typo in orthogonal init

Author: jekbradbury

jax/nn/initializers.py

@@ -94,5 +94,5 @@ def init(key, shape, dtype=np.float32):
     if n_rows < n_cols: Q = Q.T
     Q = np.reshape(Q, np.delete(shape, column_axis) + (shape[column_axis],))
     Q = np.moveaxis(Q, -1, column_axis)
-    return stddev * Q
+    return scale * Q
   return init