Repository: keras-team/keras (python)

Issue:
Question: Best practice for direct variable initialization
Some layers require explicit initialization of the whole variable tensor. What is the best practice for doing this? Currently, I am passing the tensor as the `value` argument to `keras.initializers.Constant`, but this seems like a hacky work-around, since the docs never explicitly mention using `Constant` with non-floats, and the value is also unnecessarily multiplied with a tensor of ones on initializer call.

For reference, I am doing something like this (there is no built-in initializer for rolling indices):

```py3
class SomeLayer(keras.Layer):
    def build(self, input_shape):
        indices = keras.ops.roll(keras.ops.arange(input_shape[-1]), shift=input_shape[-1] // 2)
        
        self.indices = self.add_variable(shape=input_shape[-1], initializer=keras.initializers.Constant(indices), trainable=False)
```

Pull request description:
Add support for direct tensor as initializer
Fix #19767

This PR also simplified the int8 weights assignment during quantization.

Patch:
diff --git a/keras/api/_tf_keras/keras/losses/__init__.py b/keras/api/_tf_keras/keras/losses/__init__.py
index 832d78f5f..e64b91b30 100644
--- a/keras/api/_tf_keras/keras/losses/__init__.py
+++ b/keras/api/_tf_keras/keras/losses/__init__.py
@@ -15,6 +15,7 @@ from keras.src.losses.losses import BinaryFocalCrossentropy
 from keras.src.losses.losses import CategoricalCrossentropy
 from keras.src.losses.losses import CategoricalFocalCrossentropy
 from keras.src.losses.losses import CategoricalHinge
+from keras.src.losses.losses import Circle
 from keras.src.losses.losses import CosineSimilarity
 from keras.src.losses.losses import Dice
 from keras.src.losses.losses import Hinge
@@ -34,6 +35,7 @@ from keras.src.losses.losses import binary_focal_crossentropy
 from keras.src.losses.losses import categorical_crossentropy
 from keras.src.losses.losses import categorical_focal_crossentropy
 from keras.src.losses.losses import categorical_hinge
+from keras.src.losses.losses import circle
 from keras.src.losses.losses import cosine_similarity
 from keras.src.losses.losses import ctc
 from keras.src.losses.losses import dice
diff --git a/keras/api/losses/__init__.py b/keras/api/losses/__init__.py
index ecaadddf6..af88a721c 100644
--- a/keras/api/losses/__init__.py
+++ b/keras/api/losses/__init__.py
@@ -14,6 +14,7 @@ from keras.src.losses.losses import BinaryFocalCrossentropy
 from keras.src.losses.losses import CategoricalCrossentropy
 from keras.src.losses.losses import CategoricalFocalCrossentropy
 from keras.src.losses.losses import CategoricalHinge
+from keras.src.losses.losses import Circle
 from keras.src.losses.losses import CosineSimilarity
 from keras.src.losses.losses import Dice
 from keras.src.losses.losses import Hinge
@@ -33,6 +34,7 @@ from keras.src.losses.losses import binary_focal_crossentropy
 from keras.src.losses.losses import categorical_crossentropy
 from keras.src.losses.losses import categorical_focal_crossentropy
 from keras.src.losses.losses import categorical_hinge
+from keras.src.losses.losses import circle
 from keras.src.losses.losses import cosine_similarity
 from keras.src.losses.losses import ctc
 from keras.src.losses.losses import dice
diff --git a/keras/src/backend/jax/math.py b/keras/src/backend/jax/math.py
index f359f7cfd..6b04f58a4 100644
--- a/keras/src/backend/jax/math.py
+++ b/keras/src/backend/jax/math.py
@@ -52,11 +52,7 @@ def in_top_k(targets, predictions, k):
 
 
 def logsumexp(x, axis=None, keepdims=False):
-    max_x = jnp.max(x, axis=axis, keepdims=True)
-    result = (
-        jnp.log(jnp.sum(jnp.exp(x - max_x), axis=axis, keepdims=True)) + max_x
-    )
-    return jnp.squeeze(result) if not keepdims else result
+    return jax.scipy.special.logsumexp(x, axis=axis, keepdims=keepdims)
 
 
 def qr(x, mode="reduced"):
diff --git a/keras/src/backend/numpy/math.py b/keras/src/backend/numpy/math.py
index b96fbece8..bec628f91 100644
--- a/keras/src/backend/numpy/math.py
+++ b/keras/src/backend/numpy/math.py
@@ -76,9 +76,7 @@ def in_top_k(targets, predictions, k):
 
 
 def logsumexp(x, axis=None, keepdims=False):
-    max_x = np.max(x, axis=axis, keepdims=True)
-    result = np.log(np.sum(np.exp(x - max_x), axis=axis, keepdims=True)) + max_x
-    return np.squeeze(result) if not keepdims else result
+    return scipy.special.logsumexp(x, axis=axis, keepdims=keepdims)
 
 
 def qr(x, mode="reduced"):
diff --git a/keras/src/backend/tensorflow/trainer.py b/keras/src/backend/tensorflow/trainer.py
index 98055e6cf..b375ffb2d 100644
--- a/keras/src/backend/tensorflow/trainer.py
+++ b/keras/src/backend/tensorflow/trainer.py
@@ -23,6 +23,11 @@ class TensorFlowTrainer(base_trainer.Trainer):
         self.test_function = None
         self.predict_function = None
 
+        # Specifies how many steps of the step_per_execution loop to unroll.
+        # Increasing this value can reduce kernel launch overhead,
+        # but will increase memory usage and compilation time.
+        self.unrolled_steps_per_execution = 1
+
         # Model must be created under scope of DistStrat it will be trained
         # with.
         if tf.distribute.has_strategy():
@@ -133,11 +138,31 @@ class TensorFlowTrainer(base_trainer.Trainer):
                     next_optional_inputs.has_value(),
                 )
 
-            def body(execution_step, optional_outputs, next_optional_inputs):
-                next_optional_outputs = tf.experimental.Optional.from_value(
-                    one_step_on_data(next_optional_inputs.get_value())
+            def inner_body(
+                execution_step, optional_outputs, next_optional_inputs
+            ):
+                def has_next():
+                    next_optional_outputs = tf.experimental.Optional.from_value(
+                        one_step_on_data(next_optional_inputs.get_value())
+                    )
+                    empty_outputs._element_spec = (
+                        next_optional_outputs.element_spec
+                    )
+                    return next_optional_outputs
+
+                def no_has_next():
+                    optional_outputs._element_spec = empty_outputs._element_spec
+                    return optional_outputs
+
+                next_optional_outputs = tf.cond(
+                    tf.logical_and(
+                        tf.less(execution_step, self.steps_per_execution),
+                        next_optional_inputs.has_value(),
+                    ),
+                    has_next,
+                    no_has_next,
                 )
-                empty_outputs._element_spec = next_optional_outputs.element_spec
+
                 return (
                     execution_step + 1,
                     next_optional_outputs,
@@ -150,6 +175,23 @@ class TensorFlowTrainer(base_trainer.Trainer):
                     ),
                 )
 
+            def body(execution_step, optional_outputs, next_optional_inputs):
+                for _ in range(
+                    min(
+                        self.unrolled_steps_per_execution,
+                        self.steps_per_execution,
+                    )
+                ):
+                    execution_step, optional_outputs, next_optional_inputs = (
+                        inner_body(
+                            execution_step,
+                            optional_outputs,
+                            next_optional_inputs,
+                        )
+                    )
+
+                return (execution_step, optional_outputs, next_optional_inputs)
+
             execution_step = tf.constant(0)
             next_optional_inputs = iterator.get_next_as_optional()
 
diff --git a/keras/src/backend/torch/math.py b/keras/src/backend/torch/math.py
index 2ddb26165..20e06b271 100644
--- a/keras/src/backend/torch/math.py
+++ b/keras/src/backend/torch/math.py
@@ -81,16 +81,8 @@ def in_top_k(targets, predictions, k):
 
 def logsumexp(x, axis=None, keepdims=False):
     x = convert_to_tensor(x)
-    if axis is None:
-        max_x = torch.max(x)
-        return torch.log(torch.sum(torch.exp(x - max_x))) + max_x
-
-    max_x = torch.amax(x, dim=axis, keepdim=True)
-    result = (
-        torch.log(torch.sum(torch.exp(x - max_x), dim=axis, keepdim=True))
-        + max_x
-    )
-    return torch.squeeze(result, dim=axis) if not keepdims else result
+    axis = tuple(range(x.dim())) if axis is None else axis
+    return torch.logsumexp(x, dim=axis, keepdim=keepdims)
 
 
 def qr(x, mode="reduced"):
diff --git a/keras/src/callbacks/tensorboard.py b/keras/src/callbacks/tensorboard.py
index 51831aa91..eef8b15a2 100644
--- a/keras/src/callbacks/tensorboard.py
+++ b/keras/src/callbacks/tensorboard.py
@@ -194,6 +194,7 @@ class TensorBoard(Callback):
 
         self._init_profile_batch(profile_batch)
         self._global_train_batch = 0
+        self._global_test_batch = 0
         self._previous_epoch_iterations = 0
         self._train_accumulated_time = 0
         self._batch_start_time = 0
@@ -212,11 +213,7 @@ class TensorBoard(Callback):
         self._log_write_dir = self.log_dir
 
         self._train_dir = os.path.join(self._log_write_dir, "train")
-        self._train_step = 0
-
         self._val_dir = os.path.join(self._log_write_dir, "validation")
-        self._val_step = 0
-
         self._writers = {}  # Resets writers.
 
         self._should_write_train_graph = False
@@ -409,7 +406,7 @@ class TensorBoard(Callback):
     def on_train_begin(self, logs=None):
         self._global_train_batch = 0
         self._previous_epoch_iterations = 0
-        self._push_writer(self._train_writer, self._train_step)
+        self._push_writer(self._train_writer, self._global_train_batch)
 
     def on_train_end(self, logs=None):
         self._pop_writer()
@@ -420,7 +417,7 @@ class TensorBoard(Callback):
         self._close_writers()
 
     def on_test_begin(self, logs=None):
-        self._push_writer(self._val_writer, self._val_step)
+        self._push_writer(self._val_writer, self._global_test_batch)
 
     def on_test_end(self, logs=None):
         if self.model.optimizer and hasattr(self.model.optimizer, "iterations"):
@@ -433,11 +430,6 @@ class TensorBoard(Callback):
                     )
         self._pop_writer()
 
-    def _implements_train_batch_hooks(self):
-        # Only call batch hooks when tracing or write_steps_per_second are
-        # enabled
-        return self._should_trace or self.write_steps_per_second
-
     def on_train_batch_begin(self, batch, logs=None):
         self._global_train_batch += 1
         if self.write_steps_per_second:
@@ -461,14 +453,14 @@ class TensorBoard(Callback):
             self.summary.scalar(
                 "batch_steps_per_second",
                 1.0 / batch_run_time,
-                step=self._train_step,
+                step=self._global_train_batch,
             )
 
         # `logs` isn't necessarily always a dict
         if isinstance(logs, dict):
             for name, value in logs.items():
                 self.summary.scalar(
-                    "batch_" + name, value, step=self._train_step
+                    "batch_" + name, value, step=self._global_train_batch
                 )
 
         if not self._should_trace:
@@ -481,6 +473,9 @@ class TensorBoard(Callback):
             if self._global_train_batch >= self._stop_batch:
                 self._stop_trace()
 
+    def on_test_batch_begin(self, batch, logs=None):
+        self._global_test_batch += 1
+
     def on_epoch_begin(self, epoch, logs=None):
         # Keeps track of epoch for profiling.
         if self.write_steps_per_second:
diff --git a/keras/src/initializers/__init__.py b/keras/src/initializers/__init__.py
index d3208c470..fc943cbee 100644
--- a/keras/src/initializers/__init__.py
+++ b/keras/src/initializers/__init__.py
@@ -1,5 +1,9 @@
 import inspect
 
+import numpy as np
+
+from keras.src import backend
+from keras.src import ops
 from keras.src.api_export import keras_export
 from keras.src.initializers.constant_initializers import Constant
 from keras.src.initializers.constant_initializers import Identity
@@ -82,7 +86,7 @@ def get(identifier):
     (case-sensitively).
 
     >>> identifier = 'Ones'
-    >>> keras.initializers.deserialize(identifier)
+    >>> keras.initializers.get(identifier)
     <...keras.initializers.initializers.Ones...>
 
     You can also specify `config` of the initializer to this function by passing
@@ -90,15 +94,34 @@ def get(identifier):
     the `class_name` must map to a `Initializer` class.
 
     >>> cfg = {'class_name': 'Ones', 'config': {}}
-    >>> keras.initializers.deserialize(cfg)
+    >>> keras.initializers.get(cfg)
     <...keras.initializers.initializers.Ones...>
 
     In the case that the `identifier` is a class, this method will return a new
     instance of the class by its constructor.
 
+    You may also pass a callable function with a signature that includes `shape`
+    and `dtype=None` as an identifier.
+
+    >>> fn = lambda shape, dtype=None: ops.ones(shape, dtype)
+    >>> keras.initializers.get(fn)
+    <function <lambda> at ...>
+
+    Alternatively, you can pass a backend tensor or numpy array as the
+    `identifier` to define the initializer values directly. Note that when
+    calling the initializer, the specified `shape` argument must be the same as
+    the shape of the tensor.
+
+    >>> tensor = ops.ones(shape=(5, 5))
+    >>> keras.initializers.get(tensor)
+    <function get.<locals>.initialize_fn at ...>
+
     Args:
-        identifier: String or dict that contains the initializer name or
-            configurations.
+        identifier: A string, dict, callable function, or tensor specifying
+            the initializer. If a string, it should be the name of an
+            initializer. If a dict, it should contain the configuration of an
+            initializer. Callable functions or predefined tensors are also
+            accepted.
 
     Returns:
         Initializer instance base on the input identifier.
@@ -110,6 +133,22 @@ def get(identifier):
     elif isinstance(identifier, str):
         config = {"class_name": str(identifier), "config": {}}
         obj = deserialize(config)
+    elif ops.is_tensor(identifier) or isinstance(
+        identifier, (np.generic, np.ndarray)
+    ):
+
+        def initialize_fn(shape, dtype=None):
+            dtype = backend.standardize_dtype(dtype)
+            if backend.standardize_shape(shape) != backend.standardize_shape(
+                identifier.shape
+            ):
+                raise ValueError(
+                    f"Expected `shape` to be {identifier.shape} for direct "
+                    f"tensor as initializer. Received shape={shape}"
+                )
+            return ops.cast(identifier, dtype)
+
+        obj = initialize_fn
     else:
         obj = identifier
 
diff --git a/keras/src/layers/attention/multi_head_attention.py b/keras/src/layers/attention/multi_head_attention.py
index 4dad064aa..cddfae38f 100644
--- a/keras/src/layers/attention/multi_head_attention.py
+++ b/keras/src/layers/attention/multi_head_attention.py
@@ -474,18 +474,18 @@ class MultiHeadAttention(Layer):
         #   N = `num_attention_heads`
         #   H = `size_per_head`
         # `query` = [B, T, N ,H]
-        query = self._query_dense(query)
+        query = self._query_dense.call(query)
 
         # `key` = [B, S, N, H]
-        key = self._key_dense(key)
+        key = self._key_dense.call(key)
 
         # `value` = [B, S, N, H]
-        value = self._value_dense(value)
+        value = self._value_dense.call(value)
 
         attention_output, attention_scores = self._compute_attention(
             query, key, value, attention_mask, training
         )
-        attention_output = self._output_dense(attention_output)
+        attention_output = self._output_dense.call(attention_output)
 
         if return_attention_scores:
             return attention_output, attention_scores
diff --git a/keras/src/layers/core/dense.py b/keras/src/layers/core/dense.py
index d4f1df4f4..21063a382 100644
--- a/keras/src/layers/core/dense.py
+++ b/keras/src/layers/core/dense.py
@@ -524,11 +524,7 @@ class Dense(Layer):
             del self._kernel
             # Utilize a lambda expression as an initializer to prevent adding a
             # large constant to the computation graph.
-            self._int8_build(
-                kernel_shape,
-                lambda shape, dtype: kernel_value,
-                lambda shape, dtype: kernel_scale,
-            )
+            self._int8_build(kernel_shape, kernel_value, kernel_scale)
         elif mode == "float8":
             self._float8_build()
         else:
diff --git a/keras/src/layers/core/einsum_dense.py b/keras/src/layers/core/einsum_dense.py
index 90c7b9ca3..1600ae59b 100644
--- a/keras/src/layers/core/einsum_dense.py
+++ b/keras/src/layers/core/einsum_dense.py
@@ -662,11 +662,7 @@ class EinsumDense(Layer):
             del self._kernel
             # Utilize a lambda expression as an initializer to prevent adding a
             # large constant to the computation graph.
-            self._int8_build(
-                kernel_shape,
-                lambda shape, dtype: kernel_value,
-                lambda shape, dtype: kernel_scale,
-            )
+            self._int8_build(kernel_shape, kernel_value, kernel_scale)
         elif mode == "float8":
             self._float8_build()
         else:
diff --git a/keras/src/layers/core/embedding.py b/keras/src/layers/core/embedding.py
index 6600d7209..38ced7194 100644
--- a/keras/src/layers/core/embedding.py
+++ b/keras/src/layers/core/embedding.py
@@ -358,10 +358,7 @@ class Embedding(Layer):
             del self._embeddings
             # Utilize a lambda expression as an initializer to prevent adding a
             # large constant to the computation graph.
-            self._int8_build(
-                lambda shape, dtype: embeddings_value,
-                lambda shape, dtype: embeddings_scale,
-            )
+            self._int8_build(embeddings_value, embeddings_scale)
         else:
             raise self._quantization_mode_error(mode)
 
diff --git a/keras/src/losses/__init__.py b/keras/src/losses/__init__.py
index 3163f43d9..7afeb55a0 100644
--- a/keras/src/losses/__init__.py
+++ b/keras/src/losses/__init__.py
@@ -8,6 +8,7 @@ from keras.src.losses.losses import BinaryFocalCrossentropy
 from keras.src.losses.losses import CategoricalCrossentropy
 from keras.src.losses.losses import CategoricalFocalCrossentropy
 from keras.src.losses.losses import CategoricalHinge
+from keras.src.losses.losses import Circle
 from keras.src.losses.losses import CosineSimilarity
 from keras.src.losses.losses import Dice
 from keras.src.losses.losses import Hinge
@@ -28,6 +29,7 @@ from keras.src.losses.losses import binary_focal_crossentropy
 from keras.src.losses.losses import categorical_crossentropy
 from keras.src.losses.losses import categorical_focal_crossentropy
 from keras.src.losses.losses import categorical_hinge
+from keras.src.losses.losses import circle
 from keras.src.losses.losses import cosine_similarity
 from keras.src.losses.losses import ctc
 from keras.src.losses.losses import dice
@@ -72,6 +74,8 @@ ALL_OBJECTS = {
     # Image segmentation
     Dice,
     Tversky,
+    # Similarity
+    Circle,
     # Sequence
     CTC,
     # Probabilistic
@@ -97,6 +101,8 @@ ALL_OBJECTS = {
     # Image segmentation
     dice,
     tversky,
+    # Similarity
+    circle,
     # Sequence
     ctc,
 }
diff --git a/keras/src/losses/losses.py b/keras/src/losses/losses.py
index 409c8284f..0c8ba9dfb 100644
--- a/keras/src/losses/losses.py
+++ b/keras/src/losses/losses.py
@@ -7,6 +7,7 @@ from keras.src.api_export import keras_export
 from keras.src.losses.loss import Loss
 from keras.src.losses.loss import squeeze_or_expand_to_same_rank
 from keras.src.saving import serialization_lib
+from keras.src.utils.numerical_utils import build_pos_neg_masks
 from keras.src.utils.numerical_utils import normalize
 
 
@@ -1403,6 +1404,98 @@ class Tversky(LossFunctionWrapper):
         return config
 
 
+@keras_export("keras.losses.Circle")
+class Circle(LossFunctionWrapper):
+    """Computes Circle Loss between integer labels and L2-normalized embeddings.
+
+    This is a metric learning loss designed to minimize within-class distance
+    and maximize between-class distance in a flexible manner by dynamically
+    adjusting the penalty strength based on optimization status of each
+    similarity score.
+
+    To use Circle Loss effectively, the model should output embeddings without
+    an activation function (such as a `Dense` layer with `activation=None`)
+    followed by UnitNormalization layer to ensure unit-norm embeddings.
+
+    Args:
+        gamma: Scaling factor that determines the largest scale of each
+            similarity score. Defaults to `80`.
+        margin: The relaxation factor, below this distance, negatives are
+        up weighted and positives are down weighted. Similarly, above this
+        distance negatives are down weighted and positive are up weighted.
+            Defaults to `0.4`.
+        remove_diagonal: Boolean, whether to remove self-similarities from the
+            positive mask. Defaults to `True`.
+        reduction: Type of reduction to apply to the loss. In almost all cases
+            this should be `"sum_over_batch_size"`. Supported options are
+            `"sum"`, `"sum_over_batch_size"`, `"mean"`,
+            `"mean_with_sample_weight"` or `None`. `"sum"` sums the loss,
+            `"sum_over_batch_size"` and `"mean"` sum the loss and divide by the
+            sample size, and `"mean_with_sample_weight"` sums the loss and
+            divides by the sum of the sample weights. `"none"` and `None`
+            perform no aggregation. Defaults to `"sum_over_batch_size"`.
+        name: Optional name for the loss instance.
+        dtype: The dtype of the loss's computations. Defaults to `None`, which
+            means using `keras.backend.floatx()`. `keras.backend.floatx()` is a
+            `"float32"` unless set to different value
+            (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
+            provided, then the `compute_dtype` will be utilized.
+
+    Examples:
+
+    Usage with the `compile()` API:
+
+    ```python
+    model = models.Sequential([
+        keras.layers.Input(shape=(224, 224, 3)),
+        keras.layers.Conv2D(16, (3, 3), activation='relu'),
+        keras.layers.Flatten(),
+        keras.layers.Dense(64, activation=None),  # No activation
+        keras.layers.UnitNormalization()  # L2 normalization
+    ])
+
+    model.compile(optimizer="adam", loss=keras.losses.Circle())
+    ```
+
+    Reference:
+    - [Yifan Sun et al., 2020](https://arxiv.org/abs/2002.10857)
+
+    """
+
+    def __init__(
+        self,
+        gamma=80.0,
+        margin=0.4,
+        remove_diagonal=True,
+        reduction="sum_over_batch_size",
+        name="circle",
+        dtype=None,
+    ):
+        super().__init__(
+            circle,
+            name=name,
+            reduction=reduction,
+            dtype=dtype,
+            gamma=gamma,
+            margin=margin,
+            remove_diagonal=remove_diagonal,
+        )
+        self.gamma = gamma
+        self.margin = margin
+        self.remove_diagonal = remove_diagonal
+
+    def get_config(self):
+        config = Loss.get_config(self)
+        config.update(
+            {
+                "gamma": self.gamma,
+                "margin": self.margin,
+                "remove_diagonal": self.remove_diagonal,
+            }
+        )
+        return config
+
+
 def convert_binary_labels_to_hinge(y_true):
     """Converts binary labels into -1/1 for hinge loss/metric calculation."""
     are_zeros = ops.equal(y_true, 0)
@@ -2406,3 +2499,91 @@ def tversky(y_true, y_pred, alpha=0.5, beta=0.5):
     )
 
     return 1 - tversky
+
+
+@keras_export("keras.losses.circle")
+def circle(
+    y_true,
+    y_pred,
+    ref_labels=None,
+    ref_embeddings=None,
+    remove_diagonal=True,
+    gamma=80,
+    margin=0.4,
+):
+    """Computes the Circle loss.
+
+    It is designed to minimize within-class distances and maximize between-class
+    distances in L2 normalized embedding space.
+
+    Args:
+        y_true: Tensor with ground truth labels in integer format.
+        y_pred: Tensor with predicted L2 normalized embeddings.
+        ref_labels: Optional integer tensor with labels for reference
+            embeddings. If `None`, defaults to `y_true`.
+        ref_embeddings: Optional tensor with L2 normalized reference embeddings.
+            If `None`, defaults to `y_pred`.
+        remove_diagonal: Boolean, whether to remove self-similarities from
+            positive mask. Defaults to `True`.
+        gamma: Float, scaling factor for the loss. Defaults to `80`.
+        margin: Float, relaxation factor for the loss. Defaults to `0.4`.
+
+    Returns:
+        Circle loss value.
+    """
+    y_pred = ops.convert_to_tensor(y_pred)
+    y_true = ops.cast(y_true, "int32")
+    ref_embeddings = (
+        y_pred
+        if ref_embeddings is None
+        else ops.convert_to_tensor(ref_embeddings)
+    )
+    ref_labels = y_true if ref_labels is None else ops.cast(ref_labels, "int32")
+
+    optim_pos = margin
+    optim_neg = 1 + margin
+    delta_pos = margin
+    delta_neg = 1 - margin
+
+    pairwise_cosine_distances = 1 - ops.matmul(
+        y_pred, ops.transpose(ref_embeddings)
+    )
+
+    pairwise_cosine_distances = ops.maximum(pairwise_cosine_distances, 0.0)
+    positive_mask, negative_mask = build_pos_neg_masks(
+        y_true,
+        ref_labels,
+        remove_diagonal=remove_diagonal,
+    )
+    positive_mask = ops.cast(
+        positive_mask, dtype=pairwise_cosine_distances.dtype
+    )
+    negative_mask = ops.cast(
+        negative_mask, dtype=pairwise_cosine_distances.dtype
+    )
+
+    pos_weights = optim_pos + pairwise_cosine_distances
+    pos_weights = pos_weights * positive_mask
+    pos_weights = ops.maximum(pos_weights, 0.0)
+    neg_weights = optim_neg - pairwise_cosine_distances
+    neg_weights = neg_weights * negative_mask
+    neg_weights = ops.maximum(neg_weights, 0.0)
+
+    pos_dists = delta_pos - pairwise_cosine_distances
+    neg_dists = delta_neg - pairwise_cosine_distances
+
+    pos_wdists = -1 * gamma * pos_weights * pos_dists
+    neg_wdists = gamma * neg_weights * neg_dists
+
+    p_loss = ops.logsumexp(
+        ops.where(positive_mask, pos_wdists, float("-inf")),
+        axis=1,
+    )
+    n_loss = ops.logsumexp(
+        ops.where(negative_mask, neg_wdists, float("-inf")),
+        axis=1,
+    )
+
+    circle_loss = ops.softplus(p_loss + n_loss)
+    backend.set_keras_mask(circle_loss, circle_loss > 0)
+    return circle_loss
diff --git a/keras/src/models/functional.py b/keras/src/models/functional.py
index 0acb84ed0..9c71308a6 100644
--- a/keras/src/models/functional.py
+++ b/keras/src/models/functional.py
@@ -214,8 +214,12 @@ class Functional(Function, Model):
 
     def _maybe_warn_inputs_struct_mismatch(self, inputs):
         try:
+            # We first normalize to tuples before performing the check to
+            # suppress warnings when encountering mismatched tuples and lists.
             tree.assert_same_structure(
-                inputs, self._inputs_struct, check_types=False
+                tree.lists_to_tuples(inputs),
+                tree.lists_to_tuples(self._inputs_struct),
+                check_types=False,
             )
         except:
             model_inputs_struct = tree.map_structure(
diff --git a/keras/src/utils/numerical_utils.py b/keras/src/utils/numerical_utils.py
index 0b8427551..dcd2cc422 100644
--- a/keras/src/utils/numerical_utils.py
+++ b/keras/src/utils/numerical_utils.py
@@ -193,3 +193,33 @@ def encode_categorical_inputs(
             axis=reduction_axis,
         )
         return outputs
+
+
+def build_pos_neg_masks(
+    query_labels,
+    key_labels,
+    remove_diagonal=True,
+):
+    from keras.src import ops
+
+    if ops.ndim(query_labels) == 1:
+        query_labels = ops.reshape(query_labels, (-1, 1))
+
+    if ops.ndim(key_labels) == 1:
+        key_labels = ops.reshape(key_labels, (-1, 1))
+
+    positive_mask = ops.equal(query_labels, ops.transpose(key_labels))
+    negative_mask = ops.logical_not(positive_mask)
+
+    if remove_diagonal:
+        positive_mask = ops.logical_and(
+            positive_mask,
+            ~ops.eye(
+                ops.size(query_labels),
+                ops.size(key_labels),
+                k=0,
+                dtype="bool",
+            ),
+        )
+
+    return positive_mask, negative_mask
diff --git a/keras/src/utils/torch_utils.py b/keras/src/utils/torch_utils.py
index c37044c24..ceed2425e 100644
--- a/keras/src/utils/torch_utils.py
+++ b/keras/src/utils/torch_utils.py
@@ -33,11 +33,12 @@ class TorchModuleWrapper(Layer):
     PyTorch modules.
 
     ```python
+    import torch
     import torch.nn as nn
     import torch.nn.functional as F
 
     import keras
-    from keras.src.layers import TorchModuleWrapper
+    from keras.layers import TorchModuleWrapper
 
     class Classifier(keras.Model):
         def __init__(self, **kwargs):


Test patch:
diff --git a/keras/src/initializers/random_initializers_test.py b/keras/src/initializers/random_initializers_test.py
index 06e3cadd1..a8c33a7c4 100644
--- a/keras/src/initializers/random_initializers_test.py
+++ b/keras/src/initializers/random_initializers_test.py
@@ -162,6 +162,25 @@ class InitializersTest(testing.TestCase):
         with self.assertRaises(ValueError):
             initializers.get("typo")
 
+    def test_get_method_with_tensor(self):
+        shape = (5, 5)
+
+        # Test backend tensor
+        tensor = random.uniform(shape=shape)
+        initializer = initializers.get(tensor)
+        values = initializer(shape=shape)
+        self.assertAllClose(values, tensor)
+
+        # Test numpy array
+        tensor = np.random.uniform(size=shape).astype("float32")
+        initializer = initializers.get(tensor)
+        values = initializer(shape=shape)
+        self.assertAllClose(values, tensor)
+
+        # Test bad `shape` argument
+        with self.assertRaisesRegex(ValueError, r"Expected `shape` to be"):
+            initializer(shape=(10, 10))
+
     def test_variance_scaling_invalid_scale(self):
         seed = 1234
 
diff --git a/keras/src/losses/losses_test.py b/keras/src/losses/losses_test.py
index 2ed7920a0..bbecbc06d 100644
--- a/keras/src/losses/losses_test.py
+++ b/keras/src/losses/losses_test.py
@@ -1645,3 +1645,99 @@ class TverskyTest(testing.TestCase):
         y_pred = np.array(([[4, 1], [6, 1]]))
         output = losses.Tversky(dtype="bfloat16")(y_true, y_pred)
         self.assertDType(output, "bfloat16")
+
+
+class CircleTest(testing.TestCase):
+    def setup(self):
+        self.y_true = np.array([1, 1, 2, 2, 3])
+        self.y_pred = np.array(
+            [
+                [0.70014004, -0.42008403, 0.14002801, 0.56011203],
+                [0.17609018, 0.70436073, -0.52827054, 0.44022545],
+                [-0.34050261, 0.25537696, -0.68100522, 0.59587957],
+                [0.32163376, -0.75047877, 0.53605627, -0.21442251],
+                [0.51261459, -0.34174306, 0.17087153, 0.76892189],
+            ]
+        )
+        self.ref_labels = np.array([1, 1, 2, 2, 3, 4])
+        self.ref_embeddings = np.array(
+            [
+                [0.40824829, -0.54433105, 0.27216553, 0.68041382],
+                [0.76376261, 0.10910895, -0.54554473, 0.32732684],
+                [-0.74420841, 0.24806947, 0.49613894, -0.3721042],
+                [0.52981294, -0.13245324, 0.79471941, -0.26490647],
+                [0.54554473, -0.32732684, 0.10910895, 0.76376261],
+                [-0.27216553, 0.68041382, 0.40824829, -0.54433105],
+            ]
+        )
+
+    def test_config(self):
+        self.run_class_serialization_test(
+            losses.Circle(name="mycircle", gamma=80.0, margin=0.4)
+        )
+
+    def test_correctness(self):
+        self.setup()
+        circle_loss = losses.Circle(gamma=80.0, margin=0.4)
+        loss = circle_loss(self.y_true, self.y_pred)
+        self.assertAlmostEqual(loss, 188.3883)
+
+        circle_loss = losses.Circle(gamma=256, margin=0.25)
+        loss = circle_loss(self.y_true, self.y_pred)
+        self.assertAlmostEqual(loss, 652.7617)
+
+        loss = losses.circle(
+            self.y_true,
+            self.y_pred,
+            ref_labels=self.ref_labels,
+            ref_embeddings=self.ref_embeddings,
+            gamma=80.0,
+            margin=0.4,
+            remove_diagonal=False,
+        )
+
+        self.assertAllClose(
+            loss, (61.5844, 94.3465, 276.9344, 90.9873, 48.8963)
+        )
+
+    def test_correctness_weighted(self):
+        self.setup()
+        sample_weight = np.array([2.0, 2.0, 1.0, 1.0, 0.5])
+        circle_loss = losses.Circle(gamma=80.0, margin=0.4)
+        loss = circle_loss(
+            self.y_true, self.y_pred, sample_weight=sample_weight
+        )
+        self.assertAlmostEqual(loss, 244.91918)
+
+    def test_no_reduction(self):
+        self.setup()
+        circle_loss = losses.Circle(gamma=80.0, margin=0.4, reduction=None)
+        loss = circle_loss(self.ref_labels, self.ref_embeddings)
+
+        self.assertAllClose(
+            loss, [82.9116, 36.7942, 92.4590, 52.6798, 0.0, 0.0]
+        )
+
+    def test_sum_reduction(self):
+        self.setup()
+        circle_loss = losses.Circle(gamma=80.0, margin=0.4, reduction="sum")
+        loss = circle_loss(self.ref_labels, self.ref_embeddings)
+
+        self.assertAlmostEqual(loss, 264.845)
+
+    def test_mean_with_sample_weight_reduction(self):
+        self.setup()
+        sample_weight = np.array([2.0, 2.0, 1.0, 1.0, 0.5])
+        circle_loss = losses.Circle(
+            gamma=80.0, margin=0.4, reduction="mean_with_sample_weight"
+        )
+        loss = circle_loss(
+            self.y_true, self.y_pred, sample_weight=sample_weight
+        )
+        self.assertAlmostEqual(loss, 163.27948)
+
+    def test_dtype_arg(self):
+        self.setup()
+        circle_loss = losses.Circle(dtype="bfloat16")
+        loss = circle_loss(self.y_true, self.y_pred)
+        self.assertDType(loss, "bfloat16")
diff --git a/keras/src/models/functional_test.py b/keras/src/models/functional_test.py
index ef792fd7f..44858d338 100644
--- a/keras/src/models/functional_test.py
+++ b/keras/src/models/functional_test.py
@@ -1,4 +1,5 @@
 import os
+import warnings
 
 import numpy as np
 import pytest
@@ -503,13 +504,19 @@ class FunctionalTest(testing.TestCase):
         model = Model({"i1": i1, "i2": i2}, outputs)
 
         with pytest.warns() as record:
-            model([np.ones((2, 2)), np.zeros((2, 2))])
+            model.predict([np.ones((2, 2)), np.zeros((2, 2))], verbose=0)
         self.assertLen(record, 1)
         self.assertStartsWith(
             str(record[0].message),
             r"The structure of `inputs` doesn't match the expected structure:",
         )
 
+        # No warning for mismatched tuples and lists.
+        model = Model([i1, i2], outputs)
+        with warnings.catch_warnings(record=True) as warning_logs:
+            model.predict((np.ones((2, 2)), np.zeros((2, 2))), verbose=0)
+            self.assertLen(warning_logs, 0)
+
     def test_for_functional_in_sequential(self):
         # Test for a v3.4.1 regression.
         if backend.image_data_format() == "channels_first":
diff --git a/keras/src/ops/nn_test.py b/keras/src/ops/nn_test.py
index 3a891a597..248b4c61f 100644
--- a/keras/src/ops/nn_test.py
+++ b/keras/src/ops/nn_test.py
@@ -2449,6 +2449,9 @@ class NNOpsDtypeTest(testing.TestCase):
         import jax.nn as jnn
         import jax.numpy as jnp
 
+        if dtype == "bfloat16":
+            self.skipTest("Weirdness with numpy")
+
         x = knp.ones((2), dtype=dtype)
         x_jax = jnp.ones((2), dtype=dtype)
         expected_dtype = standardize_dtype(jnn.glu(x_jax).dtype)
diff --git a/keras/src/ops/numpy_test.py b/keras/src/ops/numpy_test.py
index 1509f6102..14eef4445 100644
--- a/keras/src/ops/numpy_test.py
+++ b/keras/src/ops/numpy_test.py
@@ -3423,9 +3423,9 @@ class NumpyOneInputOpsCorrectnessTest(testing.TestCase):
         self.assertAllClose(knp.Ceil()(x), np.ceil(x))
 
     def test_clip(self):
-        x = np.array([[1.2, 2.1, -2.5], [2.4, -11.9, -5.5]])
-        self.assertAllClose(knp.clip(x, -2, 2), np.clip(x, -2, 2))
-        self.assertAllClose(knp.clip(x, -2, 2), np.clip(x, -2, 2))
+        x = np.array([[1.2, 2.1, 0.5], [2.4, 11.9, 0.5]])
+        self.assertAllClose(knp.clip(x, 1, 2), np.clip(x, 1, 2))
+        self.assertAllClose(knp.clip(x, 1, 2), np.clip(x, 1, 2))
 
         self.assertAllClose(knp.Clip(0, 1)(x), np.clip(x, 0, 1))
         self.assertAllClose(knp.Clip(0, 1)(x), np.clip(x, 0, 1))
@@ -6066,15 +6066,15 @@ class NumpyDtypeTest(testing.TestCase):
 
         x = knp.ones((1,), dtype=dtype)
         x_jax = jnp.ones((1,), dtype=dtype)
-        expected_dtype = standardize_dtype(jnp.clip(x_jax, -2, 2).dtype)
+        expected_dtype = standardize_dtype(jnp.clip(x_jax, 1, 2).dtype)
         if dtype == "bool":
             expected_dtype = "int32"
 
         self.assertEqual(
-            standardize_dtype(knp.clip(x, -2, 2).dtype), expected_dtype
+            standardize_dtype(knp.clip(x, 1, 2).dtype), expected_dtype
         )
         self.assertEqual(
-            standardize_dtype(knp.Clip(-2, 2).symbolic_call(x).dtype),
+            standardize_dtype(knp.Clip(1, 2).symbolic_call(x).dtype),
             expected_dtype,
         )
 
diff --git a/keras/src/trainers/trainer_test.py b/keras/src/trainers/trainer_test.py
index 0b19321ec..edf1555be 100644
--- a/keras/src/trainers/trainer_test.py
+++ b/keras/src/trainers/trainer_test.py
@@ -968,6 +968,75 @@ class TestTrainer(testing.TestCase):
         )
         self.assertAllClose(model.evaluate(x, y), model_2.evaluate(x, y))
 
+    @parameterized.named_parameters(
+        named_product(steps_per_execution=[3, 8, 32])
+    )
+    @pytest.mark.requires_trainable_backend
+    @pytest.mark.skipif(
+        backend.backend() != "tensorflow",
+        reason="`unrolled_steps_per_execution` is only "
+        "available with the tensorflow backend.",
+    )
+    def test_steps_per_execution_unrolled_steps_steps_count(
+        self, steps_per_execution
+    ):
+        data_size = 100
+        batch_size = 16
+        epochs = 2
+        unrolled_steps_per_execution = 8
+
+        batches_indices = list(
+            range(0, data_size, steps_per_execution * batch_size)
+        )
+
+        x = np.ones((data_size, 4))
+        y = np.ones((data_size, 1))
+
+        model = ExampleModel(units=1)
+        model.compile(
+            loss="mse",
+            optimizer="sgd",
+            steps_per_execution=steps_per_execution,
+            jit_compile=True,
+        )
+        step_count = StepCount(batches_indices, batch_size)
+        model.unrolled_steps_per_execution = unrolled_steps_per_execution
+        history = model.fit(
+            x=x,
+            y=y,
+            batch_size=batch_size,
+            epochs=epochs,
+            callbacks=[step_count],
+            verbose=0,
+        )
+
+        self.assertEqual(step_count.begin_count, len(batches_indices))
+        self.assertEqual(step_count.end_count, step_count.begin_count)
+        self.assertEqual(step_count.epoch_begin_count, epochs)
+        self.assertEqual(
+            step_count.epoch_end_count, step_count.epoch_begin_count
+        )
+
+        model_2 = ExampleModel(units=1)
+        model_2.compile(
+            loss="mse",
+            optimizer="sgd",
+            steps_per_execution=steps_per_execution,
+            jit_compile=True,
+        )
+        model_2.unrolled_steps_per_execution = 1
+        history_2 = model_2.fit(
+            x=x, y=y, batch_size=batch_size, epochs=epochs, verbose=0
+        )
+
+        self.assertAllClose(history.history["loss"], history_2.history["loss"])
+        self.assertAllClose(model.get_weights(), model_2.get_weights())
+        self.assertAllClose(
+            model.predict(x, batch_size=batch_size),
+            model_2.predict(x, batch_size=batch_size),
+        )
+        self.assertAllClose(model.evaluate(x, y), model_2.evaluate(x, y))
+
     @parameterized.named_parameters(
         named_product(
             steps_per_execution=[1, 50], mode=["eager", "non_jit", "jit"]
diff --git a/keras/src/utils/numerical_utils_test.py b/keras/src/utils/numerical_utils_test.py
index 41e2f1b3b..9b9520abc 100644
--- a/keras/src/utils/numerical_utils_test.py
+++ b/keras/src/utils/numerical_utils_test.py
@@ -72,3 +72,80 @@ class TestNumericalUtils(testing.TestCase):
         out = numerical_utils.normalize(xb, axis=-1, order=order)
         self.assertTrue(backend.is_tensor(out))
         self.assertAllClose(backend.convert_to_numpy(out), expected)
+
+    def test_build_pos_neg_masks(self):
+        query_labels = np.array([0, 1, 2, 2, 0])
+        key_labels = np.array([0, 1, 2, 0, 2])
+        expected_shape = (len(query_labels), len(key_labels))
+
+        positive_mask, negative_mask = numerical_utils.build_pos_neg_masks(
+            query_labels, key_labels, remove_diagonal=False
+        )
+
+        positive_mask = backend.convert_to_numpy(positive_mask)
+        negative_mask = backend.convert_to_numpy(negative_mask)
+        self.assertEqual(positive_mask.shape, expected_shape)
+        self.assertEqual(negative_mask.shape, expected_shape)
+        self.assertTrue(
+            np.all(np.logical_not(np.logical_and(positive_mask, negative_mask)))
+        )
+
+        expected_positive_mask_keep_diag = np.array(
+            [
+                [1, 0, 0, 1, 0],
+                [0, 1, 0, 0, 0],
+                [0, 0, 1, 0, 1],
+                [0, 0, 1, 0, 1],
+                [1, 0, 0, 1, 0],
+            ],
+            dtype="bool",
+        )
+        self.assertTrue(
+            np.all(positive_mask == expected_positive_mask_keep_diag)
+        )
+        self.assertTrue(
+            np.all(
+                negative_mask
+                == np.logical_not(expected_positive_mask_keep_diag)
+            )
+        )
+
+        positive_mask, negative_mask = numerical_utils.build_pos_neg_masks(
+            query_labels, key_labels, remove_diagonal=True
+        )
+        positive_mask = backend.convert_to_numpy(positive_mask)
+        negative_mask = backend.convert_to_numpy(negative_mask)
+        self.assertEqual(positive_mask.shape, expected_shape)
+        self.assertEqual(negative_mask.shape, expected_shape)
+        self.assertTrue(
+            np.all(np.logical_not(np.logical_and(positive_mask, negative_mask)))
+        )
+
+        expected_positive_mask_with_remove_diag = np.array(
+            [
+                [0, 0, 0, 1, 0],
+                [0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 1],
+                [0, 0, 1, 0, 1],
+                [1, 0, 0, 1, 0],
+            ],
+            dtype="bool",
+        )
+        self.assertTrue(
+            np.all(positive_mask == expected_positive_mask_with_remove_diag)
+        )
+
+        query_labels = np.array([1, 2, 3])
+        key_labels = np.array([1, 2, 3, 1])
+
+        positive_mask, negative_mask = numerical_utils.build_pos_neg_masks(
+            query_labels, key_labels, remove_diagonal=True
+        )
+        positive_mask = backend.convert_to_numpy(positive_mask)
+        negative_mask = backend.convert_to_numpy(negative_mask)
+        expected_shape_diff_sizes = (len(query_labels), len(key_labels))
+        self.assertEqual(positive_mask.shape, expected_shape_diff_sizes)
+        self.assertEqual(negative_mask.shape, expected_shape_diff_sizes)
+        self.assertTrue(
+            np.all(np.logical_not(np.logical_and(positive_mask, negative_mask)))
+        )