Introduce dtype inference and improve dtype in ops.numpy.*

keras_core/backend/jax/numpy.py

@@ -70,11 +70,13 @@
     return jnp.max(x, axis=axis, keepdims=keepdims, initial=initial)
 
 
-def ones(shape, dtype="float32"):
+def ones(shape, dtype=None):
+    dtype = dtype or config.floatx()
     return jnp.ones(shape, dtype=dtype)
 
 
-def zeros(shape, dtype="float32"):
+def zeros(shape, dtype=None):
+    dtype = dtype or config.floatx()
     return jnp.zeros(shape, dtype=dtype)
 
 
@@ -253,7 +255,8 @@
     return jnp.dot(x, y)
 
 
-def empty(shape, dtype="float32"):
+def empty(shape, dtype=None):
+    dtype = dtype or config.floatx()
     return jnp.empty(shape, dtype=dtype)
 
 
@@ -284,6 +287,7 @@
 
 
 def full(shape, fill_value, dtype=None):
+    dtype = dtype or config.floatx()
     return jnp.full(shape, fill_value, dtype=dtype)
 
 
@@ -307,7 +311,8 @@
     return jnp.hstack(xs)
 
 
-def identity(n, dtype="float32"):
+def identity(n, dtype=None):
+    dtype = dtype or config.floatx()
     return jnp.identity(n, dtype=dtype)
 
 
@@ -348,6 +353,7 @@
 def linspace(
     start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0
 ):
+    dtype = dtype or config.floatx()
     return jnp.linspace(
         start,
         stop,
@@ -398,6 +404,7 @@
 
 
 def logspace(start, stop, num=50, endpoint=True, base=10, dtype=None, axis=0):
+    dtype = dtype or config.floatx()
     return jnp.logspace(
         start,
         stop,
@@ -573,7 +580,8 @@
     return jnp.trace(x, offset=offset, axis1=axis1, axis2=axis2)
 
 
-def tri(N, M=None, k=0, dtype="float32"):
+def tri(N, M=None, k=0, dtype=None):
+    dtype = dtype or config.floatx()
     return jnp.tri(N, M=M, k=k, dtype=dtype)
 
 
@@ -652,7 +660,8 @@
     return jnp.sum(x, axis=axis, keepdims=keepdims)
 
 
-def eye(N, M=None, k=0, dtype="float32"):
+def eye(N, M=None, k=0, dtype=None):
+    dtype = dtype or config.floatx()
     return jnp.eye(N, M=M, k=k, dtype=dtype)
 
 

keras_core/backend/numpy/numpy.py

@@ -34,11 +34,13 @@
     return np.max(x, axis=axis, keepdims=keepdims, initial=initial)
 
 
-def ones(shape, dtype="float32"):
+def ones(shape, dtype=None):
+    dtype = dtype or config.floatx()
     return np.ones(shape, dtype=dtype)
 
 
-def zeros(shape, dtype="float32"):
+def zeros(shape, dtype=None):
+    dtype = dtype or config.floatx()
     return np.zeros(shape, dtype=dtype)
 
 
@@ -134,7 +136,6 @@
 
 
 def array(x, dtype=None):
-    dtype = dtype or config.floatx()
     return np.array(x, dtype=dtype)
 
 
@@ -251,7 +252,8 @@
     return np.dot(x, y)
 
 
-def empty(shape, dtype="float32"):
+def empty(shape, dtype=None):
+    dtype = dtype or config.floatx()
     return np.empty(shape, dtype=dtype)
 
 
@@ -302,7 +304,8 @@
     return np.hstack(xs)
 
 
-def identity(n, dtype="float32"):
+def identity(n, dtype=None):
+    dtype = dtype or config.floatx()
     return np.identity(n, dtype=dtype)
 
 
@@ -338,6 +341,7 @@
     start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0
 ):
     axis = tuple(axis) if isinstance(axis, list) else axis
+    dtype = dtype or config.floatx()
     return np.linspace(
         start,
         stop,
@@ -382,6 +386,7 @@
 
 
 def logspace(start, stop, num=50, endpoint=True, base=10, dtype=None, axis=0):
+    dtype = dtype or config.floatx()
     return np.logspace(
         start,
         stop,
@@ -556,7 +561,8 @@
     return np.trace(x, offset=offset, axis1=axis1, axis2=axis2)
 
 
-def tri(N, M=None, k=0, dtype="float32"):
+def tri(N, M=None, k=0, dtype=None):
+    dtype = dtype or config.floatx()
     return np.tri(N, M=M, k=k, dtype=dtype)
 
 
@@ -631,7 +637,8 @@
     return np.sum(x, axis=axis, keepdims=keepdims)
 
 
-def eye(N, M=None, k=0, dtype="float32"):
+def eye(N, M=None, k=0, dtype=None):
+    dtype = dtype or config.floatx()
     return np.eye(N, M=M, k=k, dtype=dtype)
 
 

keras_core/backend/tensorflow/numpy.py

@@ -196,11 +196,13 @@
     return tfnp.max(x, axis=axis, keepdims=keepdims)
 
 
-def ones(shape, dtype="float32"):
+def ones(shape, dtype=None):
+    dtype = dtype or config.floatx()
     return tf.ones(shape, dtype=dtype)
 
 
-def zeros(shape, dtype="float32"):
+def zeros(shape, dtype=None):
+    dtype = dtype or config.floatx()
     return tf.zeros(shape, dtype=dtype)
 
 
@@ -403,7 +405,8 @@
     return tfnp.dot(x, y)
 
 
-def empty(shape, dtype="float32"):
+def empty(shape, dtype=None):
+    dtype = dtype or config.floatx()
     return tfnp.empty(shape, dtype=dtype)
 
 
@@ -434,6 +437,7 @@
 
 
 def full(shape, fill_value, dtype=None):
+    dtype = dtype or config.floatx()
     return tfnp.full(shape, fill_value, dtype=dtype)
 
 
@@ -453,7 +457,8 @@
     return tfnp.hstack(xs)
 
 
-def identity(n, dtype="float32"):
+def identity(n, dtype=None):
+    dtype = dtype or config.floatx()
     return tfnp.identity(n, dtype=dtype)
 
 
@@ -488,6 +493,7 @@
 def linspace(
     start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0
 ):
+    dtype = dtype or config.floatx()
     return tfnp.linspace(
         start,
         stop,
@@ -532,6 +538,7 @@
 
 
 def logspace(start, stop, num=50, endpoint=True, base=10, dtype=None, axis=0):
+    dtype = dtype or config.floatx()
     return tfnp.logspace(
         start,
         stop,
@@ -776,7 +783,8 @@
     return tfnp.trace(x, offset=offset, axis1=axis1, axis2=axis2)
 
 
-def tri(N, M=None, k=0, dtype="float32"):
+def tri(N, M=None, k=0, dtype=None):
+    dtype = dtype or config.floatx()
     return tfnp.tri(N, M=M, k=k, dtype=dtype)
 
 
@@ -863,7 +871,8 @@
     return tfnp.sum(x, axis=axis, keepdims=keepdims)
 
 
-def eye(N, M=None, k=0, dtype="float32"):
+def eye(N, M=None, k=0, dtype=None):
+    dtype = dtype or config.floatx()
     return tfnp.eye(N, M=M, k=k, dtype=dtype)
 
 

keras_core/backend/torch/numpy.py

@@ -76,15 +76,15 @@
     return result
 
 
-def ones(shape, dtype="float32"):
-    dtype = to_torch_dtype(dtype)
+def ones(shape, dtype=None):
+    dtype = to_torch_dtype(dtype or config.floatx())
     if isinstance(shape, int):
         shape = (shape,)
     return torch.ones(size=shape, dtype=dtype, device=get_device())
 
 
-def zeros(shape, dtype="float32"):
-    dtype = to_torch_dtype(dtype)
+def zeros(shape, dtype=None):
+    dtype = to_torch_dtype(dtype or config.floatx())
     if isinstance(shape, int):
         shape = (shape,)
     return torch.zeros(size=shape, dtype=dtype, device=get_device())
@@ -230,7 +230,8 @@
 
 
 def array(x, dtype=None):
-    dtype = to_torch_dtype(dtype)
+    if dtype is not None:
+        dtype = to_torch_dtype(dtype)
     if isinstance(x, torch.Tensor):
         return x
     return torch.tensor(x, dtype=dtype, device=get_device())
@@ -386,8 +387,8 @@
     return torch.matmul(x, y)
 
 
-def empty(shape, dtype="float32"):
-    dtype = to_torch_dtype(dtype)
+def empty(shape, dtype=None):
+    dtype = to_torch_dtype(dtype or config.floatx())
     return torch.empty(size=shape, dtype=dtype, device=get_device())
 
 
@@ -426,7 +427,7 @@
 
 
 def full(shape, fill_value, dtype=None):
-    dtype = to_torch_dtype(dtype)
+    dtype = to_torch_dtype(dtype or config.floatx())
     fill_value = convert_to_tensor(fill_value, dtype=dtype)
     if len(fill_value.shape) > 0:
         # `torch.full` only supports scala `fill_value`.
@@ -457,8 +458,8 @@
     return torch.hstack(xs)
 
 
-def identity(n, dtype="float32"):
-    dtype = to_torch_dtype(dtype)
+def identity(n, dtype=None):
+    dtype = to_torch_dtype(dtype or config.floatx())
     return torch.eye(n, dtype=dtype)
 
 
@@ -512,7 +513,7 @@
             "torch.linspace does not support an `axis` argument. "
             f"Received axis={axis}"
         )
-    dtype = to_torch_dtype(dtype)
+    dtype = to_torch_dtype(dtype or config.floatx())
     if endpoint is False:
         stop = stop - ((stop - start) / num)
     if hasattr(start, "__len__") and hasattr(stop, "__len__"):
@@ -586,7 +587,7 @@
             "torch.logspace does not support an `axis` argument. "
             f"Received axis={axis}"
         )
-    dtype = to_torch_dtype(dtype)
+    dtype = to_torch_dtype(dtype or config.floatx())
     if endpoint is False:
         stop = stop - ((stop - start) / num)
     if hasattr(start, "__len__") and hasattr(stop, "__len__"):
@@ -738,7 +739,8 @@
 
 def prod(x, axis=None, keepdims=False, dtype=None):
     x = convert_to_tensor(x)
-    dtype = to_torch_dtype(dtype)
+    if dtype is not None:
+        dtype = to_torch_dtype(dtype)
     if axis is None:
         return torch.prod(x, dtype=dtype)
     if not isinstance(axis, (list, tuple)):
@@ -933,8 +935,8 @@
     return torch.sum(torch.diagonal(x, offset, axis1, axis2), dim=-1)
 
 
-def tri(N, M=None, k=0, dtype="float32"):
-    dtype = to_torch_dtype(dtype)
+def tri(N, M=None, k=0, dtype=None):
+    dtype = to_torch_dtype(dtype or config.floatx())
     M = M or N
     x = torch.ones((N, M), dtype=dtype, device=get_device())
     return torch.tril(x, diagonal=k)
@@ -1037,8 +1039,8 @@
     return torch.sum(x)
 
 
-def eye(N, M=None, k=None, dtype="float32"):
-    dtype = to_torch_dtype(dtype)
+def eye(N, M=None, k=None, dtype=None):
+    dtype = to_torch_dtype(dtype or config.floatx())
     M = N if M is None else M
     k = 0 if k is None else k
     if k == 0:

keras_core/layers/rnn/dropout_rnn_cell_test.py

@@ -64,4 +64,31 @@ def test_basics(self):
             expected_num_non_trainable_weights=0,
             expected_num_non_trainable_variables=1,
             supports_masking=True,
+            run_mixed_precision_check=False,
         )
+
+        # Custom mixed_float16 check
+        # Never test mixed precision on torch CPU. Torch lacks support.
+        run_mixed_precision_check = True
+        if backend.backend() == "torch":
+            import torch
+
+            run_mixed_precision_check = torch.cuda.is_available()
+        if run_mixed_precision_check:
+            self.run_layer_test(
+                layers.RNN,
+                init_kwargs={
+                    "cell": RNNCellWithDropout(
+                        5, seed=1337, dtype="mixed_float16"
+                    ),
+                    "dtype": "mixed_float16",
+                },
+                input_shape=(3, 2, 4),
+                call_kwargs={"training": True},
+                expected_output_shape=(3, 5),
+                expected_num_trainable_weights=2,
+                expected_num_non_trainable_weights=0,
+                expected_num_non_trainable_variables=1,
+                supports_masking=True,
+                run_mixed_precision_check=False,
+            )