keras/src/backend/tensorflow/optimizer.py/TFOptimizer/add_variable_from_reference
class TFOptimizer:    def add_variable_from_reference(
        self, reference_variable, name=None, initializer="zeros"
    ):
        if isinstance(reference_variable, backend.Variable):
            colocate_var = reference_variable.value
        else:
            colocate_var = reference_variable

        with self._distribution_strategy.extended.colocate_vars_with(
            colocate_var
        ):
            return super().add_variable_from_reference(
                reference_variable, name=name, initializer=initializer
            )