def METHOD_NAME(self, reference_variable, name=None):
    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().METHOD_NAME(
            reference_variable, name=name
        )