def _tf_apply(self, grads, trainable_variables=None):
        """Tensorflow specific logic for apply, which handles distribution."""
        from keras.src.utils.module_utils import tensorflow as tf

        if tf.distribute.in_cross_replica_context():
            raise ValueError("apply() must be called in a replica context.")

        if tf.__internal__.distribute.strategy_supports_no_merge_call():
            self._common_apply(grads, trainable_variables=trainable_variables)
        else:

            def _handle_cross_replica(distribution, grads, trainable_variables):
                finite_per_replica = (
                    distribution.extended.call_for_each_replica(
                        self.check_finite, args=(grads,)
                    )
                )
                # Each replica computed the same `finite` value, since
                # `grads` is all-reduced across replicas. Arbitrarily take
                # `finite` from the first replica.
                finite = distribution.experimental_local_results(
                    finite_per_replica
                )[0]

                def apply_fn():
                    distribution.extended.call_for_each_replica(
                        self._stateful_handle_finite_grads,
                        args=(grads, trainable_variables),
                    )

                # Note: We must call this cond() in a cross-replica context.
                # DistributionStrategy does not support having a cond in a
                # replica context with a branch that calls `merge_call`, and
                # self._optimizer.apply_gradients calls `merge_call`.
                ops.cond(
                    finite, apply_fn, self._stateful_handle_non_finite_grads
                )

            tf.distribute.get_replica_context().merge_call(
                _handle_cross_replica, args=(grads, trainable_variables)
            )
Tensorflow specific logic for apply, which handles distribution.
keras/src/optimizers/loss_scale_optimizer.py