Prevent accidental upcasting in jax.nn.initializers.

Currently distribution parameters such as stddev and scale are expected to be
weakly typed scalars.  When they're passed as float32 they can cause an upcast
of the initialized arrays even when the dtype is specified as e.g. bfloat16.
Some users were surprised by this.

PiperOrigin-RevId: 611858446

Author: levskaya

jax/_src/nn/initializers.py

@@ -130,7 +130,7 @@ def init(key: KeyArray,
            shape: core.Shape,
            dtype: DTypeLikeInexact = dtype) -> Array:
     dtype = dtypes.canonicalize_dtype(dtype)
-    return random.uniform(key, shape, dtype) * scale
+    return random.uniform(key, shape, dtype) * jnp.array(scale, dtype)
   return init
 
 @export
@@ -156,7 +156,7 @@ def init(key: KeyArray,
            shape: core.Shape,
            dtype: DTypeLikeInexact = dtype) -> Array:
     dtype = dtypes.canonicalize_dtype(dtype)
-    return random.normal(key, shape, dtype) * stddev
+    return random.normal(key, shape, dtype) * jnp.array(stddev, dtype)
   return init
 
 @export
@@ -193,7 +193,8 @@ def init(key: KeyArray,
            shape: core.Shape,
            dtype: DTypeLikeInexact = dtype) -> Array:
     dtype = dtypes.canonicalize_dtype(dtype)
-    return random.truncated_normal(key, lower, upper, shape, dtype) * stddev
+    return random.truncated_normal(
+        key, lower, upper, shape, dtype) * jnp.array(stddev, dtype)
   return init
 
 @export
@@ -613,7 +614,7 @@ def init(key: KeyArray,
     if n_rows < n_cols: Q = Q.T
     Q = jnp.reshape(Q, tuple(np.delete(shape, column_axis)) + (shape[column_axis],))
     Q = jnp.moveaxis(Q, -1, column_axis)
-    return scale * Q
+    return jnp.array(scale, dtype) * Q
   return init
 
 @export

tests/nn_test.py

@@ -403,6 +403,17 @@ def testVarianceScalingError(self):
     ):
       initializer(rng, shape)
 
+  def testAccidentalUpcasting(self):
+    rng = random.PRNGKey(0)
+    shape = (4, 4)
+    scalar_param = jnp.array(1.0, dtype=jnp.float32)
+    for init_fn in (nn.initializers.uniform(scalar_param, jnp.bfloat16),
+                    nn.initializers.normal(scalar_param, jnp.bfloat16),
+                    nn.initializers.truncated_normal(scalar_param, jnp.bfloat16),
+                   ):
+      sub_rng, rng = random.split(rng)
+      val = init_fn(sub_rng, shape)
+      self.assertEqual(val.dtype, jnp.bfloat16)
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())