# Copyright 2018 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.
# ==============================================================================
"""Strategy combinations for combinations.combine()."""

import sys
import unittest
from tensorflow.core.protobuf import config_pb2
from tensorflow.python import tf2
from tensorflow.python.distribute import central_storage_strategy
from tensorflow.python.distribute import cluster_resolver
from tensorflow.python.distribute import collective_all_reduce_strategy
from tensorflow.python.distribute import combinations
from tensorflow.python.distribute import distribute_lib
from tensorflow.python.distribute import mirrored_strategy as mirrored_lib
from tensorflow.python.distribute import multi_process_runner
from tensorflow.python.distribute import multi_worker_test_base
from tensorflow.python.distribute import one_device_strategy as one_device_lib
from tensorflow.python.distribute import parameter_server_strategy_v2
from tensorflow.python.distribute import sharded_variable
from tensorflow.python.distribute import test_util
from tensorflow.python.distribute import tpu_strategy as tpu_lib
from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
from tensorflow.python.eager import context
from tensorflow.python.eager import remote
from tensorflow.python.framework import device as tf_device
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util as framework_test_util
from tensorflow.python.platform import flags
from tensorflow.python.tpu import device_assignment as device_assignment_lib
from tensorflow.python.training import server_lib
from tensorflow.python.util.tf_export import tf_export

_TF_INTERNAL_API_PREFIX = "__internal__.distribute.combinations."

_did_connect_to_cluster = False
_topology = None
CollectiveAllReduceExtended = (
    collective_all_reduce_strategy.CollectiveAllReduceExtended)


def _version_chooser(tf1_cls, tf2_cls):

  def creator(*args, **kwargs):
    if tf2.enabled():
      return tf2_cls(*args, **kwargs)
    return tf1_cls(*args, **kwargs)

  return creator


MirroredStrategy = _version_chooser(mirrored_lib.MirroredStrategyV1,
                                    mirrored_lib.MirroredStrategy)
CentralStorageStrategy = _version_chooser(
    central_storage_strategy.CentralStorageStrategyV1,
    central_storage_strategy.CentralStorageStrategy)
OneDeviceStrategy = _version_chooser(one_device_lib.OneDeviceStrategyV1,
                                     one_device_lib.OneDeviceStrategy)
# Only V2 CollectiveAllReduceStrategy combinations are supported.
CollectiveAllReduceStrategy = (
    collective_all_reduce_strategy.CollectiveAllReduceStrategy)


# pylint: disable=missing-docstring
def _get_tpu_strategy_creator(steps_per_run,
                              use_single_core=False,
                              enable_packed_variable=False,
                              enable_spmd_xla_paritioning=False,
                              **kwargs):

  def _create_tpu_strategy():
    FLAGS = flags.FLAGS  # pylint: disable=invalid-name
    global _did_connect_to_cluster
    global _topology

    try:
      # Attempt to locally discover the TPU. This will fail for Cloud TPU, in
      # which case we fall back to the values passed as flags.
      resolver = tpu_cluster_resolver.TPUClusterResolver()
      did_automatically_resolve = True
    except ValueError:
      did_automatically_resolve = False

      # These flags will be defined by tpu_test_wrapper.py.
      resolver = tpu_cluster_resolver.TPUClusterResolver(
          tpu=hasattr(FLAGS, "tpu") and FLAGS.tpu or "",
          zone=hasattr(FLAGS, "zone") and FLAGS.zone or None,
          project=hasattr(FLAGS, "project") and FLAGS.project or None,
      )

    # Only connect once per process, rather than per test method.
    if not _did_connect_to_cluster:
      if getattr(FLAGS, "tpu", "") or did_automatically_resolve:
        remote.connect_to_cluster(resolver)
        _did_connect_to_cluster = True
      _topology = tpu_cluster_resolver.initialize_tpu_system(resolver)

    device_assignment = None
    if use_single_core:
      device_assignment = device_assignment_lib.DeviceAssignment(
          _topology,
          core_assignment=device_assignment_lib.SINGLE_CORE_ASSIGNMENT)

    # Steps per run is only supported in TF 1.x
    if tf2.enabled():
      strategy = tpu_lib.TPUStrategyV2(
          resolver,
          device_assignment,
          experimental_spmd_xla_partitioning=enable_spmd_xla_paritioning,
          **kwargs)
    else:
      strategy = tpu_lib.TPUStrategyV1(resolver, steps_per_run,
                                       device_assignment, **kwargs)
    if enable_packed_variable and enable_spmd_xla_paritioning:
      raise ValueError("Packed Variable is not compatiable with SPMD mode")
    strategy._enable_packed_variable_in_eager_mode = enable_packed_variable  # pylint: disable=protected-access
    return strategy

  return _create_tpu_strategy


def _mirrored_strategy_with_collective_key_base(devices):
  required_cpus_nums = sum(
      1
      for d in devices
      if tf_device.DeviceSpec.from_string(d).device_type == "CPU"
  )

  # If required virtual CPUs are not setup yet, config the logical devices.
  if required_cpus_nums > len(context.context().list_logical_devices("CPU")):
    context._reset_context()  # pylint: disable=protected-access
    test_util.set_logical_devices_to_at_least("CPU", required_cpus_nums)

  # Increase collective base key to avoid key collision across subtests.
  mirrored_lib.MirroredStrategyV1._collective_key_base += 100000
  mirrored_lib.MirroredStrategy._collective_key_base += 100000
  return MirroredStrategy(devices)


def _mirrored_strategy_with_no_merge_call(devices):
  mirrored_lib.MirroredStrategyV1._collective_key_base += 100000
  mirrored_lib.MirroredStrategy._collective_key_base += 100000
  out = MirroredStrategy(devices)
  # Stub out merge call usage.
  out.extended._use_merge_call = lambda: False  # pylint: disable=protected-access
  return out


def _get_multi_worker_mirrored_creator(required_gpus, use_merge_call=True):

  def _create_multi_worker_mirrored():
    tf_config = cluster_resolver.TFConfigClusterResolver()
    master = tf_config.master()
    if tf_config.rpc_layer:
      # Strip off the rpc_layer suffix.
      master = master[len("%s://" % tf_config.rpc_layer):]
    resolver = cluster_resolver.SimpleClusterResolver(
        cluster_spec=tf_config.cluster_spec(),
        task_type=tf_config.task_type,
        task_id=tf_config.task_id,
        master=master,
        environment=tf_config.environment,
        num_accelerators={"GPU": required_gpus},
        rpc_layer=tf_config.rpc_layer or "grpc",
    )
    # Disable health check and coordination service. We don't have a reliable
    # way to shutdown the strategy (and thus the strategy health check or
    # coordination service heartbeat) at the end of a test. Turning on the
    # strategy health check or coordination service heartbeat causes some
    # flakiness since we re-create part of the server when creating a strategy,
    # and our tests are capable of handling failures.
    CollectiveAllReduceExtended._enable_check_health = False  # pylint: disable=protected-access
    context.context().configure_coordination_service(service_type="")
    # Always create the strategy in eager mode so that it starts the server and
    # configures the eager context. The eager context can no longer be
    # configured after initialization.
    with context.eager_mode():
      strategy = CollectiveAllReduceStrategy(cluster_resolver=resolver)

    if not use_merge_call:
      strategy.extended._use_merge_call = lambda: False  # pylint: disable=protected-access
    # TODO(b/152320929): Wait for the cluster before proceeding, otherwise
    # collectives may hang if any worker launches collectives before the chief
    # creates the strategy.
    try:
      multi_process_runner.get_barrier().wait()
    except ValueError:
      # If the creator is called in the main process,
      # multi_process_runner.get_barrier() raises ValueError, which is safe to
      # ignore.
      pass
    return strategy

  def skip_if_cannot_start_grpc_server():
    try:
      return _create_multi_worker_mirrored()
    except errors.UnknownError as e:
      if "Could not start gRPC server" in e.message and (
          len(sys.argv) >= 1 and "bazel" in sys.argv[0]):
        raise unittest.SkipTest("Cannot start std servers.")
      else:
        raise

  return skip_if_cannot_start_grpc_server


# Due to b/195615322, FixedShardsPartitioner will wrongly partition
# RNG state, so we use MinSizePartitioner as the default. Maximum RNG
# state size is int64[3] which is 8 * 3 bytes, so we set
# min_shard_bytes to 8 * 3 + 1.
DEFAULT_PARTITIONER = sharded_variable.MinSizePartitioner(
    min_shard_bytes=8 * 3 + 1, max_shards=2)


def _get_ps_strategy_creator(num_workers,
                             num_ps,
                             required_gpus=0,
                             variable_partitioner=DEFAULT_PARTITIONER):

  def _create_ps_strategy(resolver, variable_partitioner):
    return parameter_server_strategy_v2.ParameterServerStrategyV2(
        resolver, variable_partitioner=variable_partitioner)

  def _create_parameter_server():
    if framework_test_util.is_xla_enabled():
      # To address test failures resulting in XLA with MultiProcessRunner,
      # continue to use in-process cluster for XLA tests.
      cluster_def = multi_worker_test_base.create_in_process_cluster(
          num_workers=num_workers, num_ps=num_ps, rpc_layer="grpc")
      resolver = cluster_resolver.SimpleClusterResolver(
          server_lib.ClusterSpec(cluster_def),
          num_accelerators={"GPU": required_gpus},
          rpc_layer="grpc")
      return _create_ps_strategy(resolver, variable_partitioner)
    else:
      tf_config = cluster_resolver.TFConfigClusterResolver()
      cluster_def = tf_config.cluster_spec().as_dict()
      if not cluster_def:
        # When MultiProcessRunner cluster is used, the cluster is not created
        # initially when the decorator is called. When the test runs, initially
        # this method is invoked via decorator before setting up the
        # MultiProcessRunner with worker and ps in the combinations.py. After
        # setup is done, the subprocess invokes this method again to get
        # strategy object. We return None strategy when the main thread invokes
        # this method before setting up cluster.
        # Returning None is fine here, since this thread will proceed to create
        # MultiProcessRunner and invoke tests with decorator inside
        # subprocesses.
        return None
      # MultiProcessRunner is already setup and this method is invoked from a
      # subprocess running the actual test.
      resolver = cluster_resolver.SimpleClusterResolver(
          server_lib.ClusterSpec(cluster_def),
          num_accelerators={"GPU": required_gpus},
          task_type=tf_config.task_type,
          task_id=tf_config.task_id,
          environment=tf_config.environment,
          rpc_layer=tf_config.rpc_layer or "grpc")
      if tf_config.task_type in ("worker", "ps"):
        worker_config = config_pb2.ConfigProto()
        worker_config.inter_op_parallelism_threads = 4  # max num_workers + 1

        try:
          server = server_lib.Server(
              cluster_def,
              job_name=tf_config.task_type,
              task_index=tf_config.task_id,
              protocol="grpc",
              config=worker_config)
        except errors.UnknownError as e:
          if "Could not start gRPC server" in e.message:
            raise unittest.SkipTest("Cannot start std servers.")
          else:
            raise

        # Blocking the process that starts a server from exiting.
        server.join()

      return _create_ps_strategy(resolver, variable_partitioner)

  return _create_parameter_server


def _deferred_pool_runner(has_chief,
                          num_workers,
                          initializer=None,
                          share_gpu=True):
  """Returns a callable that returns the pool runner.

  It creates the pool runner only upon first invocation. This avoids creating it
  when this file is imported.

  Args:
    has_chief: whether there should be a chief.
    num_workers: the number of workers excluding the chief.
    initializer: initializer of each process.
    share_gpu: whether to share GPU between the workers.

  Returns:
    A callable that returns the runner.
  """

  container = []

  def get_or_create():
    if not container:
      cluster_spec = multi_worker_test_base.create_cluster_spec(
          has_chief=has_chief,
          num_workers=num_workers,
          num_ps=0,
          has_eval=False)
      runner = multi_process_runner.MultiProcessPoolRunner(
          cluster_spec, initializer=initializer, share_gpu=share_gpu)
      container.append(runner)
    return container[0]

  return get_or_create


# We need to create the strategy in the initializer to start the server before
# any test runs.
_two_worker_pool = _deferred_pool_runner(
    has_chief=True,
    num_workers=1,
    initializer=_get_multi_worker_mirrored_creator(required_gpus=0))

# Two-worker pool where each worker gets it's own GPU. Useful for testing MWMS
# on a single host.
_two_worker_pool_noshare = _deferred_pool_runner(
    has_chief=True,
    num_workers=1,
    initializer=_get_multi_worker_mirrored_creator(required_gpus=0),
    share_gpu=False)
_four_worker_pool = _deferred_pool_runner(
    has_chief=True,
    num_workers=3,
    initializer=_get_multi_worker_mirrored_creator(required_gpus=0))

# pylint: disable=g-long-lambda
default_strategy = combinations.NamedDistribution(
    "Default",
    distribute_lib._get_default_strategy,  # pylint: disable=protected-access
    required_gpus=None)
one_device_strategy = combinations.NamedDistribution(
    "OneDeviceCPU", lambda: OneDeviceStrategy("/cpu:0"), required_gpus=None)
one_device_strategy_gpu = combinations.NamedDistribution(
    "OneDeviceGPU", lambda: OneDeviceStrategy("/gpu:0"), required_gpus=1)
one_device_strategy_on_worker_1 = combinations.NamedDistribution(
    "OneDeviceOnWorker1CPU",
    lambda: OneDeviceStrategy("/job:worker/replica:0/task:1/cpu:0"),
    required_gpus=None)
one_device_strategy_gpu_on_worker_1 = combinations.NamedDistribution(
    "OneDeviceOnWorker1GPU",
    lambda: OneDeviceStrategy("/job:worker/replica:0/task:1/gpu:0"),
    required_gpus=1)
tpu_strategy = combinations.NamedDistribution(
    "TPU", _get_tpu_strategy_creator(steps_per_run=2), required_tpu=True)
tpu_strategy_packed_var = combinations.NamedDistribution(
    "TPUPackedVar",
    _get_tpu_strategy_creator(steps_per_run=2, enable_packed_variable=True),
    required_tpu=True)
tpu_strategy_spmd = combinations.NamedDistribution(
    "TPUUseSPMD",
    _get_tpu_strategy_creator(
        steps_per_run=2, enable_spmd_xla_paritioning=True),
    required_tpu=True)
tpu_strategy_one_step = combinations.NamedDistribution(
    "TPUOneStep", _get_tpu_strategy_creator(steps_per_run=1), required_tpu=True)
tpu_strategy_one_core = combinations.NamedDistribution(
    "TPUOneCore",
    _get_tpu_strategy_creator(steps_per_run=2, use_single_core=True),
    required_tpu=True)
tpu_strategy_one_step_one_core = combinations.NamedDistribution(
    "TPUOneStepOneCore",
    _get_tpu_strategy_creator(steps_per_run=1, use_single_core=True),
    required_tpu=True)
cloud_tpu_strategy = combinations.NamedDistribution(
    "CloudTPU",
    _get_tpu_strategy_creator(steps_per_run=2),
    required_tpu=True,
    use_cloud_tpu=True)
mirrored_strategy_with_one_cpu = combinations.NamedDistribution(
    "Mirrored1CPU",
    lambda: _mirrored_strategy_with_collective_key_base(["/cpu:0"]))
mirrored_strategy_with_one_gpu = combinations.NamedDistribution(
    "Mirrored1GPU",
    lambda: _mirrored_strategy_with_collective_key_base(["/gpu:0"]),
    required_gpus=1)
mirrored_strategy_with_gpu_and_cpu = combinations.NamedDistribution(
    "MirroredCPUAndGPU",
    lambda: _mirrored_strategy_with_collective_key_base(["/gpu:0", "/cpu:0"]),
    required_gpus=1)
mirrored_strategy_with_two_cpus = combinations.NamedDistribution(
    "Mirrored2CPUs",
    lambda: _mirrored_strategy_with_collective_key_base(["/cpu:0", "/cpu:1"]),
    required_gpus=0)
mirrored_strategy_with_two_gpus = combinations.NamedDistribution(
    "Mirrored2GPUs",
    lambda: _mirrored_strategy_with_collective_key_base(["/gpu:0", "/gpu:1"]),
    required_gpus=2)
mirrored_strategy_with_two_gpus_no_merge_call = combinations.NamedDistribution(
    "Mirrored2GPUsNoMergeCall",
    lambda: _mirrored_strategy_with_no_merge_call(["/gpu:0", "/gpu:1"]),
    required_physical_gpus=2)
# Should call set_virtual_cpus_to_at_least(3) in your test's setUp methods.
# Deprecated, use mirrored_strategy_with_two_cpus instead.
mirrored_strategy_with_cpu_1_and_2 = combinations.NamedDistribution(
    "Mirrored2CPU",
    lambda: _mirrored_strategy_with_collective_key_base(["/cpu:1", "/cpu:2"]))
mirrored_strategy_with_cpu_1_and_2.__doc__ = (
    """Mirrored strategy with 2 virtual CPUs.

    Should set up logical devices before use
    """)
central_storage_strategy_with_two_gpus = combinations.NamedDistribution(
    "CentralStorage2GPUs",
    lambda: CentralStorageStrategy(["/gpu:0", "/gpu:1"]),
    required_gpus=2)
central_storage_strategy_with_gpu_and_cpu = combinations.NamedDistribution(
    "CentralStorageCPUAndGPU",
    lambda: CentralStorageStrategy(["/gpu:0", "/cpu:0"]),
    required_gpus=1)
# chief + 1 worker, with CPU.
multi_worker_mirrored_2x1_cpu = combinations.NamedDistribution(
    "MultiWorkerMirrored2x1CPU",
    _get_multi_worker_mirrored_creator(required_gpus=0),
    has_chief=True,
    num_workers=1,
    pool_runner_fn=_two_worker_pool,
    no_xla=True,
)
# chief + 1 worker, with 1 GPU each.
multi_worker_mirrored_2x1_gpu = combinations.NamedDistribution(
    "MultiWorkerMirrored2x1GPU",
    _get_multi_worker_mirrored_creator(required_gpus=1),
    has_chief=True,
    num_workers=1,
    required_gpus=1,
    pool_runner_fn=_two_worker_pool,
    share_gpu=False,
)

# Same as above, but not sharing the GPU between the workers.
multi_worker_mirrored_2x1_gpu_noshare = combinations.NamedDistribution(
    "MultiWorkerMirrored2x1GPUNoShare",
    _get_multi_worker_mirrored_creator(required_gpus=1),
    has_chief=True,
    num_workers=1,
    required_gpus=1,
    pool_runner_fn=_two_worker_pool_noshare,
    share_gpu=False,
)
# chief + 1 worker, with 2 GPU each.
multi_worker_mirrored_2x2_gpu = combinations.NamedDistribution(
    "MultiWorkerMirrored2x2GPU",
    _get_multi_worker_mirrored_creator(required_gpus=2),
    has_chief=True,
    num_workers=1,
    required_gpus=2,
    pool_runner_fn=_two_worker_pool,
    no_xla=True,
)
multi_worker_mirrored_2x2_gpu_no_merge_call = combinations.NamedDistribution(
    "MultiWorkerMirrored2x2GPUNoMergeCall",
    _get_multi_worker_mirrored_creator(required_gpus=2, use_merge_call=False),
    has_chief=True,
    num_workers=1,
    required_physical_gpus=2,
    pool_runner_fn=_two_worker_pool,
    no_xla=True,
)
# chief + 3 workers, with CPU.
multi_worker_mirrored_4x1_cpu = combinations.NamedDistribution(
    "MultiWorkerMirrored4x1CPU",
    _get_multi_worker_mirrored_creator(required_gpus=0),
    has_chief=True,
    num_workers=3,
    pool_runner_fn=_four_worker_pool,
    no_xla=True,
)


def parameter_server_strategy_fn(name,
                                 num_workers,
                                 num_ps,
                                 required_gpus=0,
                                 variable_partitioner=DEFAULT_PARTITIONER):
  return combinations.NamedDistribution(
      name,
      _get_ps_strategy_creator(
          num_workers=num_workers,
          num_ps=num_ps,
          required_gpus=required_gpus,
          variable_partitioner=variable_partitioner),
      required_gpus=required_gpus,
      num_workers=num_workers,
      has_chief=True,
      num_ps=num_ps)


parameter_server_strategy_3worker_2ps_cpu = parameter_server_strategy_fn(
    "ParameterServer3Worker2PSCPU", num_workers=3, num_ps=2)
parameter_server_strategy_1worker_2ps_cpu = parameter_server_strategy_fn(
    "ParameterServer1Worker2PSCPU", num_workers=1, num_ps=2)
parameter_server_strategy_3worker_2ps_1gpu = parameter_server_strategy_fn(
    "ParameterServer3Worker2PS1GPU", num_workers=3, num_ps=2, required_gpus=1)
parameter_server_strategy_1worker_2ps_1gpu = parameter_server_strategy_fn(
    "ParameterServer1Worker2PS1GPU", num_workers=1, num_ps=2, required_gpus=1)

graph_and_eager_modes = ["graph", "eager"]


# TODO(crccw): remove after tf-nightly picks up the new API.
def set_virtual_cpus_to_at_least(num_virtual_cpus):
  test_util.set_logical_devices_to_at_least("CPU", num_virtual_cpus)


strategies_minus_tpu = [
    default_strategy,
    one_device_strategy,
    one_device_strategy_gpu,
    mirrored_strategy_with_gpu_and_cpu,
    mirrored_strategy_with_two_gpus,
    central_storage_strategy_with_gpu_and_cpu,
]

strategies_minus_default_and_tpu = [
    one_device_strategy,
    one_device_strategy_gpu,
    mirrored_strategy_with_gpu_and_cpu,
    mirrored_strategy_with_two_gpus,
]

tpu_strategies = [
    tpu_strategy,  # steps_per_run=2
    tpu_strategy_one_step,
    tpu_strategy_packed_var,
    cloud_tpu_strategy,
]

all_strategies_minus_default = strategies_minus_default_and_tpu + tpu_strategies

all_strategies = strategies_minus_tpu + tpu_strategies

two_replica_strategies = [
    mirrored_strategy_with_gpu_and_cpu,
    mirrored_strategy_with_two_gpus,
    multi_worker_mirrored_2x1_cpu,
    multi_worker_mirrored_2x1_gpu,
    tpu_strategy,  # steps_per_run=2
    tpu_strategy_one_step,
    central_storage_strategy_with_gpu_and_cpu,
]

four_replica_strategies = [
    multi_worker_mirrored_2x2_gpu,
    multi_worker_mirrored_4x1_cpu,
]

# TODO(b/159831907): replace with two_replica_strategies after the tests using
# it work with MWMS.
multidevice_strategies = [
    mirrored_strategy_with_gpu_and_cpu,
    mirrored_strategy_with_two_gpus,
    tpu_strategy,  # steps_per_run=2
    tpu_strategy_one_step
]

multiworker_strategies = [
    multi_worker_mirrored_2x1_cpu, multi_worker_mirrored_2x1_gpu,
    multi_worker_mirrored_2x2_gpu
]


def strategy_minus_tpu_combinations():
  return combinations.combine(
      distribution=strategies_minus_tpu, mode=["graph", "eager"])


def tpu_strategy_combinations():
  return combinations.combine(distribution=tpu_strategies, mode=["graph"])


def all_strategy_combinations():
  return strategy_minus_tpu_combinations() + tpu_strategy_combinations()


def all_strategy_minus_default_and_tpu_combinations():
  return combinations.combine(
      distribution=[
          one_device_strategy, one_device_strategy_gpu,
          mirrored_strategy_with_gpu_and_cpu, mirrored_strategy_with_two_gpus
      ],
      mode=["graph", "eager"])


def all_strategy_combinations_minus_default():
  return (all_strategy_minus_default_and_tpu_combinations() +
          tpu_strategy_combinations())


tf_export(
    _TF_INTERNAL_API_PREFIX + "central_storage_strategy_with_gpu_and_cpu",
    v1=[]).export_constant(__name__,
                           "central_storage_strategy_with_gpu_and_cpu")
tf_export(
    _TF_INTERNAL_API_PREFIX + "central_storage_strategy_with_two_gpus",
    v1=[]).export_constant(__name__, "central_storage_strategy_with_two_gpus")
tf_export(
    _TF_INTERNAL_API_PREFIX + "cloud_tpu_strategy",
    v1=[]).export_constant(__name__, "cloud_tpu_strategy")
tf_export(
    _TF_INTERNAL_API_PREFIX + "default_strategy",
    v1=[]).export_constant(__name__, "default_strategy")
tf_export(
    _TF_INTERNAL_API_PREFIX + "mirrored_strategy_with_cpu_1_and_2",
    v1=[]).export_constant(__name__, "mirrored_strategy_with_cpu_1_and_2")
tf_export(
    _TF_INTERNAL_API_PREFIX + "mirrored_strategy_with_two_cpus",
    v1=[]).export_constant(__name__, "mirrored_strategy_with_two_cpus")
tf_export(
    _TF_INTERNAL_API_PREFIX + "mirrored_strategy_with_gpu_and_cpu",
    v1=[]).export_constant(__name__, "mirrored_strategy_with_gpu_and_cpu")
tf_export(
    _TF_INTERNAL_API_PREFIX + "mirrored_strategy_with_one_cpu",
    v1=[]).export_constant(__name__, "mirrored_strategy_with_one_cpu")
tf_export(
    _TF_INTERNAL_API_PREFIX + "mirrored_strategy_with_one_gpu",
    v1=[]).export_constant(__name__, "mirrored_strategy_with_one_gpu")
tf_export(
    _TF_INTERNAL_API_PREFIX + "mirrored_strategy_with_two_gpus",
    v1=[]).export_constant(__name__, "mirrored_strategy_with_two_gpus")
tf_export(
    _TF_INTERNAL_API_PREFIX + "mirrored_strategy_with_two_gpus_no_merge_call",
    v1=[]).export_constant(__name__,
                           "mirrored_strategy_with_two_gpus_no_merge_call")
tf_export(
    _TF_INTERNAL_API_PREFIX + "multi_worker_mirrored_2x1_cpu",
    v1=[]).export_constant(__name__, "multi_worker_mirrored_2x1_cpu")
tf_export(
    _TF_INTERNAL_API_PREFIX + "multi_worker_mirrored_2x1_gpu",
    v1=[]).export_constant(__name__, "multi_worker_mirrored_2x1_gpu")
tf_export(
    _TF_INTERNAL_API_PREFIX + "multi_worker_mirrored_2x1_gpu_noshare",
    v1=[]).export_constant(__name__, "multi_worker_mirrored_2x1_gpu_noshare")
tf_export(
    _TF_INTERNAL_API_PREFIX + "multi_worker_mirrored_2x2_gpu",
    v1=[]).export_constant(__name__, "multi_worker_mirrored_2x2_gpu")
tf_export(
    _TF_INTERNAL_API_PREFIX + "multi_worker_mirrored_2x2_gpu_no_merge_call",
    v1=[]).export_constant(__name__,
                           "multi_worker_mirrored_2x2_gpu_no_merge_call")
tf_export(
    _TF_INTERNAL_API_PREFIX + "one_device_strategy",
    v1=[]).export_constant(__name__, "one_device_strategy")
tf_export(
    _TF_INTERNAL_API_PREFIX + "one_device_strategy_gpu",
    v1=[]).export_constant(__name__, "one_device_strategy_gpu")
tf_export(
    _TF_INTERNAL_API_PREFIX + "tpu_strategy",
    v1=[]).export_constant(__name__, "tpu_strategy")
tf_export(
    _TF_INTERNAL_API_PREFIX + "parameter_server_strategy_3worker_2ps_cpu",
    v1=[]).export_constant(__name__,
                           "parameter_server_strategy_3worker_2ps_cpu")
tf_export(
    _TF_INTERNAL_API_PREFIX + "parameter_server_strategy_1worker_2ps_cpu",
    v1=[]).export_constant(__name__,
                           "parameter_server_strategy_1worker_2ps_cpu")
tf_export(
    _TF_INTERNAL_API_PREFIX + "parameter_server_strategy_3worker_2ps_1gpu",
    v1=[]).export_constant(__name__,
                           "parameter_server_strategy_3worker_2ps_1gpu")
tf_export(
    _TF_INTERNAL_API_PREFIX + "parameter_server_strategy_1worker_2ps_1gpu",
    v1=[]).export_constant(__name__,
                           "parameter_server_strategy_1worker_2ps_1gpu")
tf_export(
    _TF_INTERNAL_API_PREFIX + "tpu_strategy_one_core",
    v1=[]).export_constant(__name__, "tpu_strategy_one_core")
tf_export(
    _TF_INTERNAL_API_PREFIX + "tpu_strategy_packed_var",
    v1=[]).export_constant(__name__, "tpu_strategy_packed_var")