Repository: google/jax

Issue:
stax neural network initializers default to 32bit floats even in 64 bit mode
I need/want to try out my code with 64 bit wide floats. My code uses neural networks setup with stax, the initializers for which default to 32 bit floats (and there is no useable API to change this in user facing code).

Using the thus initialized (`float32`) parameters (and `float64` data batches) in an update `fori_loop` (taking the gradients with `value_and_grad` and updating parameters with any optimizer) results in either one of two faulty behaviors (depending on the exact nature of computation):

1. the gradients come out as `float32`, thus thwarting my intention of computing with `float64` values
2. the gradients come out as `float64`, which results in a type-mismatch exception thrown by `fori_loop` since the initial value argument passed to it was of type `float32`




Discussion:
`stax` is intentionally a somewhat minimal example of a neural network library that is perhaps more intended for expository purposes than anything else. Its goal is to be simple more than it is to be feature complete.

There are a number of other neural network libraries built on top of JAX that may be more useful and more feature complete (examples include [flax](https://github.com/google-research/flax/tree/prerelease), [trax](https://github.com/google/trax), [Jaxnet](https://github.com/JuliusKunze/jaxnet), ... and there are more! A thousand flowers seem to be blooming right now, which is great!) I don't know if any of them in particular support 64-bit weights, but it seems like a reasonable thing to want.

That said, `stax` is very minimal, only a few hundred lines of code. So if you are mostly happy with stax, I'd actually just suggest forking it into your repository and making whatever changes you like. Or feel free to send a PR allowing more dtype flexibility. Up to you. It might also be possible to get the type behavior you want by passing in a custom initializer.
Would you consider a PR that just changes all default `dtypes` in `jax/nn/initializers.py` to `np.float64`? With default jax behavior (i.e., 32 bit mode) those will be cast to `np.float32` anyways so nothing changes there but in 64 bit mode (i.e., when `jax_enable_x64` was explicitly set to `True`) they would then always be `float64`.

I don't think that would be a good idea, because our plan is to make x64 mode the default sometime soon. The proposed change would mean that the default is simply 64-bit weights, which seems like a strange default for neural network training. I think the best plan is to make the dtype user-selectable somehow.

Resolving patch:
diff --git a/jax/nn/initializers.py b/jax/nn/initializers.py
--- a/jax/nn/initializers.py
+++ b/jax/nn/initializers.py
@@ -32,13 +32,13 @@
 def zeros(key, shape, dtype=np.float32): return np.zeros(shape, dtype)
 def ones(key, shape, dtype=np.float32): return np.ones(shape, dtype)
 
-def uniform(scale=1e-2):
-  def init(key, shape, dtype=np.float32):
+def uniform(scale=1e-2, dtype=np.float32):
+  def init(key, shape, dtype=dtype):
     return random.uniform(key, shape, dtype) * scale
   return init
 
-def normal(stddev=1e-2):
-  def init(key, shape, dtype=np.float32):
+def normal(stddev=1e-2, dtype=np.float32):
+  def init(key, shape, dtype=dtype):
     return random.normal(key, shape, dtype) * stddev
   return init
 
@@ -48,8 +48,8 @@ def _compute_fans(shape, in_axis=-2, out_axis=-1):
   fan_out = shape[out_axis] * receptive_field_size
   return fan_in, fan_out
 
-def variance_scaling(scale, mode, distribution, in_axis=-2, out_axis=-1):
-  def init(key, shape, dtype=np.float32):
+def variance_scaling(scale, mode, distribution, in_axis=-2, out_axis=-1, dtype=np.float32):
+  def init(key, shape, dtype=dtype):
     fan_in, fan_out = _compute_fans(shape, in_axis, out_axis)
     if mode == "fan_in": denominator = fan_in
     elif mode == "fan_out": denominator = fan_out
@@ -77,14 +77,14 @@ def init(key, shape, dtype=np.float32):
 kaiming_uniform = he_uniform = partial(variance_scaling, 2.0, "fan_in", "uniform")
 kaiming_normal = he_normal = partial(variance_scaling, 2.0, "fan_in", "truncated_normal")
 
-def orthogonal(scale=1.0, column_axis=-1):
+def orthogonal(scale=1.0, column_axis=-1, dtype=np.float32):
   """
   Construct an initializer for uniformly distributed orthogonal matrices.
   
   If the shape is not square, the matrices will have orthonormal rows or columns
   depending on which side is smaller.
   """
-  def init(key, shape, dtype=np.float32):
+  def init(key, shape, dtype=dtype):
     if len(shape) < 2:
       raise ValueError("orthogonal initializer requires at least a 2D shape")
     n_rows, n_cols = onp.prod(shape) // shape[column_axis], shape[column_axis]
@@ -99,21 +99,21 @@ def init(key, shape, dtype=np.float32):
   return init
 
 
-def delta_orthogonal(scale=1.0, column_axis=-1):
+def delta_orthogonal(scale=1.0, column_axis=-1, dtype=np.float32):
   """
   Construct an initializer for delta orthogonal kernels; see arXiv:1806.05393. 
 
   The shape must be 3D, 4D or 5D.
   """
-  def init(key, shape, dtype=np.float32):
+  def init(key, shape, dtype=dtype):
     if len(shape) not in [3, 4, 5]:
       raise ValueError("Delta orthogonal initializer requires a 3D, 4D or 5D "
                        "shape.")
     if shape[-1] < shape[-2]:
       raise ValueError("`fan_in` must be less or equal than `fan_out`. ")
-    ortho_init = orthogonal(scale=scale, column_axis=column_axis)
+    ortho_init = orthogonal(scale=scale, column_axis=column_axis, dtype=dtype)
     ortho_matrix = ortho_init(key, shape[-2:])
-    W = np.zeros(shape)
+    W = np.zeros(shape, dtype=dtype)
     if len(shape) == 3:
       k = shape[0]
       return ops.index_update(W, ops.index[(k-1)//2, ...], ortho_matrix)


Test patch:
diff --git a/tests/nn_test.py b/tests/nn_test.py
--- a/tests/nn_test.py
+++ b/tests/nn_test.py
@@ -74,16 +74,16 @@ def initializer_record(name, initializer, min_dims=2, max_dims=4):
   return InitializerRecord(name, initializer, shapes)
 
 INITIALIZER_RECS = [
-    initializer_record("uniform", nn.initializers.uniform(), 1),
-    initializer_record("normal", nn.initializers.normal(), 1),
-    initializer_record("he_normal", nn.initializers.he_normal()),
-    initializer_record("he_uniform", nn.initializers.he_uniform()),
-    initializer_record("glorot_normal", nn.initializers.glorot_normal()),
-    initializer_record("glorot_uniform", nn.initializers.glorot_uniform()),
-    initializer_record("lecun_normal", nn.initializers.lecun_normal()),
-    initializer_record("lecun_uniform", nn.initializers.lecun_uniform()),
-    initializer_record("orthogonal", nn.initializers.orthogonal(), 2, 2),
-    initializer_record("orthogonal", nn.initializers.delta_orthogonal(), 4, 4)
+    initializer_record("uniform", nn.initializers.uniform, 1),
+    initializer_record("normal", nn.initializers.normal, 1),
+    initializer_record("he_normal", nn.initializers.he_normal),
+    initializer_record("he_uniform", nn.initializers.he_uniform),
+    initializer_record("glorot_normal", nn.initializers.glorot_normal),
+    initializer_record("glorot_uniform", nn.initializers.glorot_uniform),
+    initializer_record("lecun_normal", nn.initializers.lecun_normal),
+    initializer_record("lecun_uniform", nn.initializers.lecun_uniform),
+    initializer_record("orthogonal", nn.initializers.orthogonal, 2, 2),
+    initializer_record("delta_orthogonal", nn.initializers.delta_orthogonal, 4, 4)
 ]
 
 class NNInitializersTest(jtu.JaxTestCase):
@@ -93,7 +93,7 @@ class NNInitializersTest(jtu.JaxTestCase):
        "_{}_{}".format(
            rec.name,
            jtu.format_shape_dtype_string(shape, dtype)),
-       "initializer": rec.initializer,
+       "initializer": rec.initializer(),
        "shape": shape, "dtype": dtype}
       for rec in INITIALIZER_RECS
       for shape in rec.shapes
@@ -101,6 +101,27 @@ class NNInitializersTest(jtu.JaxTestCase):
   def testInitializer(self, initializer, shape, dtype):
     rng = random.PRNGKey(0)
     val = initializer(rng, shape, dtype)
+    self.assertEqual(shape, np.shape(val))
+    self.assertEqual(jax.dtypes.canonicalize_dtype(dtype), np.dtype(val))
+
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name":
+       "_{}_{}".format(
+           rec.name,
+           jtu.format_shape_dtype_string(shape, dtype)),
+       "initializer_provider": rec.initializer,
+       "shape": shape, "dtype": dtype}
+      for rec in INITIALIZER_RECS
+      for shape in rec.shapes
+      for dtype in [onp.float32, onp.float64]))
+  def testInitializerProvider(self, initializer_provider, shape, dtype):
+    rng = random.PRNGKey(0)
+    initializer = initializer_provider(dtype=dtype)
+    val = initializer(rng, shape)
+
+    self.assertEqual(shape, np.shape(val))
+    self.assertEqual(jax.dtypes.canonicalize_dtype(dtype), np.dtype(val))
+
 
 if __name__ == "__main__":
   absltest.main()