Introduce DTypePolicyMap

Introduce `DTypePolicyMap`

Fix `LayerNormalization.get_config` (keras-team#19807)

Propagate kwargs through `keras.ops.isclose` (keras-team#19782)

* propagate kwargs through isclose
this allows passing atol and rtol

* switch isclose **kwargs to explicit kwargs

* reduce line lengths

* fix ops.isclose signature

* fix ops.IsClose compute_output_spec signature

* implement isclose rtol atol equal_nan args for all backends

* shorten line lengths again

* revert using tf.experimental.numpy.isclose
tensorflow version now uses code inspired from tf.experimental.numpy.isclose

* fix lint

* add docs for new parameters

Faster in_top_k implementation for Jax backend (keras-team#19814)

* Faster in_top_k implementation.

* Fix bug in rank computation.

Fix CI

Fix TypeError in `Lambda.from_config` (keras-team#19827)

fixing dmtree.is_nested() and parameterized tree test (keras-team#19822)

Fix `keras.ops.repeat` cannot return an expected shape when `x` is a … (keras-team#19826)

* Fix `keras.ops.repeat` cannot return an expected shape when `x` is a `KerasTensor` and the `axis` is `None`

* Test dynamic is still dynamic after repetition

* Improve error messages

`Metric.variables` is now recursive. (keras-team#19830)

This allows it to surface variables from metrics nested at any depth.

Previously, metrics within metrics within metrics would not have their variables tracked in JAX, causing them to not be updated.

Fix `get_file` when the HTTP response has no `Content-Length` header (keras-team#19833)

Add `ops.switch` (keras-team#19834)

* Add `ops.switch`

* Update tests

* Fix out-of-bound issue

* Revert `torch.cond`

Use `absl.testing.parameterized` for `tree_test.py`. (keras-team#19842)

For consistency, use `absl.testing.parameterized` instead of `parameterized` for `tree_test.py` since that is used for all other tests.

It's one less dependency. It also says `optree` or `dmtree` in each test name.

Make batch norm mask shape error more descriptive (keras-team#19829)

* Made batch norm mask shape error more descriptive

* Added shape info in mask error message to help with degugging

Fix code style

doc: `ops.slice` (keras-team#19843)

corrected the example code in unit_normalization.py (keras-team#19845)

Added missing closing bracket and exact output value in example code after replicating the code.

Adjust code example

Add `training` argument to `Model.compute_loss()`. (keras-team#19840)

This allows models to perform different computations during training and evaluation. For instance, some expensive to compute metrics can be skipped during training and only computed during evaluation.

Note that backwards compatibility with overrides that do not have the `training` argument is maintained.

Fix the compatibility issues of `Orthogonal` and `GRU` (keras-team#19844)

* Add legacy `Orthogonal` class name

* Add legacy `implementation` arg to `GRU`

Fix inconsistent behavior of `losses.sparse_categorical_crossentropy`… (keras-team#19838)

* Fix inconsistent behavior of `losses.sparse_categorical_crossentropy` with and without `ignore_class`

* Test

* chore(format)

* Fix tests in `losses`

Fix bugs with `Mean`, `Accuracy` and `BinaryAccuracy` metrics. (keras-team#19847)

- `reduce_to_samplewise_values` would not reduce `sample_weights` correctly because the number of dimensions of `values` was checked.
- `reduce_to_samplewise_values` needs to explicitely broadcast `sample_weights`. Before, it was implicitly broadcast in the multiplication with `values`. However, the explicit broadcast is needed for the computation of `num_samples` for the averaging to be correct. This causes a bug when `sample_weights` is of rank 2 or more and a broadcast happens when doing the multiplication. This logic existed in `tf_keras`: https://github.com/keras-team/tf-keras/blob/master/tf_keras/metrics/base_metric.py#L508
- `Accuracy` and `BinaryAccuracy` were doing a mean reduction too early, before multiplying by `sample_weights`. This matters when the rank of `sample_weights` is the same as `y_true` and `y_pred`.

Add tests for `DTypePolicyMap`

Fix test

Update the logic of `default_policy`

Improve serialization of `DTypePolicyMap`

Improve `__repr__` and `__eq__`

Add `custom_gradient` for the numpy backend (keras-team#19849)

fix variable name when add in init function (keras-team#19853)

Address comments

Update docstrings

Author: Grvzard

keras/api/_tf_keras/keras/ops/__init__.py

@@ -22,6 +22,7 @@
 from keras.src.ops.core import slice
 from keras.src.ops.core import slice_update
 from keras.src.ops.core import stop_gradient
+from keras.src.ops.core import switch
 from keras.src.ops.core import unstack
 from keras.src.ops.core import vectorized_map
 from keras.src.ops.core import while_loop

keras/api/ops/__init__.py

@@ -22,6 +22,7 @@
 from keras.src.ops.core import slice
 from keras.src.ops.core import slice_update
 from keras.src.ops.core import stop_gradient
+from keras.src.ops.core import switch
 from keras.src.ops.core import unstack
 from keras.src.ops.core import vectorized_map
 from keras.src.ops.core import while_loop

keras/src/backend/jax/core.py

@@ -287,6 +287,10 @@ def slice_update(inputs, start_indices, updates):
     return jax.lax.dynamic_update_slice(inputs, updates, start_indices)
 
 
+def switch(index, branches, *operands):
+    return jax.lax.switch(index, branches, *operands)
+
+
 def while_loop(
     cond,
     body,

keras/src/backend/jax/math.py

@@ -40,11 +40,11 @@ def top_k(x, k, sorted=True):
 
 
 def in_top_k(targets, predictions, k):
-    targets = targets[..., None]
-    topk_values = top_k(predictions, k)[0]
-    targets_values = jnp.take_along_axis(predictions, targets, axis=-1)
-    mask = targets_values >= topk_values
-    return jnp.any(mask, axis=1)
+    preds_at_label = jnp.take_along_axis(
+        predictions, jnp.expand_dims(targets, axis=-1), axis=-1
+    )
+    rank = 1 + jnp.sum(jnp.greater(predictions, preds_at_label), axis=-1)
+    return jnp.less_equal(rank, k)
 
 
 def logsumexp(x, axis=None, keepdims=False):

keras/src/backend/jax/numpy.py

@@ -598,10 +598,10 @@ def imag(x):
     return jnp.imag(x)
 
 
-def isclose(x1, x2):
+def isclose(x1, x2, rtol=1e-5, atol=1e-8, equal_nan=False):
     x1 = convert_to_tensor(x1)
     x2 = convert_to_tensor(x2)
-    return jnp.isclose(x1, x2)
+    return jnp.isclose(x1, x2, rtol, atol, equal_nan)
 
 
 @sparse.densifying_unary

keras/src/backend/jax/trainer.py

@@ -63,6 +63,7 @@ def compute_loss_and_updates(
             y=y,
             y_pred=y_pred,
             sample_weight=sample_weight,
+            training=training,
         )
         if losses:
             self._losses_override.clear()

keras/src/backend/numpy/core.py

@@ -1,3 +1,5 @@
+import warnings
+
 import numpy as np
 
 from keras.src import tree
@@ -241,6 +243,12 @@ def slice_update(inputs, start_indices, updates):
     return inputs
 
 
+def switch(index, branches, *operands):
+    index = convert_to_tensor(index, "int32")
+    index = np.clip(index, 0, len(branches) - 1)
+    return branches[index](*operands)
+
+
 def while_loop(
     cond,
     body,
@@ -279,7 +287,21 @@ def unstack(x, num=None, axis=0):
     return [x[i] for i in range(x.shape[0])]
 
 
-def custom_gradient(fun):
-    raise NotImplementedError(
-        "`custom_gradient` is not supported with numpy backend"
-    )
+class custom_gradient:
+    """Decorator for custom gradients.
+
+    Args:
+        fun: Forward pass function.
+    """
+
+    def __init__(self, fun):
+        warnings.warn(
+            "`custom_gradient` for the numpy backend acts as a pass-through to "
+            "support the forward pass. No gradient computation or modification "
+            "takes place."
+        )
+        self.fun = fun
+
+    def __call__(self, *args, **kwargs):
+        outputs, _ = self.fun(*args, **kwargs)
+        return outputs

keras/src/backend/numpy/numpy.py

@@ -65,6 +65,8 @@ def matmul(x1, x2):
         dtype = "int32"
     else:
         dtype = dtypes.result_type(x1.dtype, x2.dtype)
+    x1 = x1.astype(dtype)
+    x2 = x2.astype(dtype)
     return np.matmul(x1, x2).astype(dtype)
 
 
@@ -505,8 +507,8 @@ def imag(x):
     return np.imag(x)
 
 
-def isclose(x1, x2):
-    return np.isclose(x1, x2)
+def isclose(x1, x2, rtol=1e-5, atol=1e-8, equal_nan=False):
+    return np.isclose(x1, x2, rtol, atol, equal_nan)
 
 
 def isfinite(x):

keras/src/backend/numpy/trainer.py

@@ -28,8 +28,8 @@ def test_step(self, data):
             y_pred = self(x, training=False)
         else:
             y_pred = self(x)
-        loss = self.compute_loss(
-            x=x, y=y, y_pred=y_pred, sample_weight=sample_weight
+        loss = self._compute_loss(
+            x=x, y=y, y_pred=y_pred, sample_weight=sample_weight, training=False
         )
         self._loss_tracker.update_state(
             loss, sample_weight=tree.flatten(x)[0].shape[0]

keras/src/backend/tensorflow/core.py

@@ -355,6 +355,19 @@ def slice_update(inputs, start_indices, updates):
     return dynamic_update_slice(inputs, updates, start_indices)
 
 
+def switch(index, branches, *operands):
+    index = convert_to_tensor(index, "int32")
+    index = tf.clip_by_value(index, 0, len(branches) - 1)
+
+    # Workaround to deal with python closures. More details:
+    # https://github.com/tensorflow/tensorflow/issues/8776#issuecomment-311383887
+    def gen_fn(i):
+        return lambda: branches[i](*operands)
+
+    branch_fns = [gen_fn(i) for i in range(len(branches))]
+    return tf.switch_case(index, branch_fns)
+
+
 def while_loop(
     cond,
     body,