Internal change

PiperOrigin-RevId: 383667471

official/recommendation/ranking/train.py

@@ -23,12 +23,12 @@ from absl import logging
 import orbit
 import tensorflow as tf
 
+from official.common import distribute_utils
 from official.core import base_trainer
 from official.core import train_lib
 from official.core import train_utils
 from official.recommendation.ranking import common
 from official.recommendation.ranking.task import RankingTask
-from official.utils.misc import distribution_utils
 from official.utils.misc import keras_utils
 
 FLAGS = flags.FLAGS
@@ -86,7 +86,7 @@ def main(_) -> None:
 
   enable_tensorboard = params.trainer.callbacks.enable_tensorboard
 
-  strategy = distribution_utils.get_distribution_strategy(
+  strategy = distribute_utils.get_distribution_strategy(
       distribution_strategy=params.runtime.distribution_strategy,
       all_reduce_alg=params.runtime.all_reduce_alg,
       num_gpus=params.runtime.num_gpus,

official/utils/misc/distribution_utils.py

-# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Helper functions for running models in a distributed setting."""
-# pylint: disable=wildcard-import
-from official.common.distribute_utils import *