# Copyright 2020 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 tensorflow.compat.v2 as tf

multidevice_strategies = [
    tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu,
    tf.__internal__.distribute.combinations.mirrored_strategy_with_two_gpus,
    tf.__internal__.distribute.combinations.tpu_strategy,
]

multiworker_strategies = [
    tf.__internal__.distribute.combinations.multi_worker_mirrored_2x1_cpu,
    tf.__internal__.distribute.combinations.multi_worker_mirrored_2x1_gpu,
    tf.__internal__.distribute.combinations.multi_worker_mirrored_2x2_gpu,
]

strategies_minus_default_minus_tpu = [
    tf.__internal__.distribute.combinations.one_device_strategy,
    tf.__internal__.distribute.combinations.one_device_strategy_gpu,
    tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu,
    tf.__internal__.distribute.combinations.mirrored_strategy_with_two_gpus,
    tf.__internal__.distribute.combinations.central_storage_strategy_with_gpu_and_cpu,  # noqa: E501
]

strategies_minus_tpu = [
    tf.__internal__.distribute.combinations.default_strategy,
    tf.__internal__.distribute.combinations.one_device_strategy,
    tf.__internal__.distribute.combinations.one_device_strategy_gpu,
    tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu,
    tf.__internal__.distribute.combinations.mirrored_strategy_with_two_gpus,
    tf.__internal__.distribute.combinations.central_storage_strategy_with_gpu_and_cpu,  # noqa: E501
]

multi_worker_mirrored_strategies = [
    tf.__internal__.distribute.combinations.multi_worker_mirrored_2x1_cpu,
    tf.__internal__.distribute.combinations.multi_worker_mirrored_2x1_gpu,
    tf.__internal__.distribute.combinations.multi_worker_mirrored_2x2_gpu,
]

tpu_strategies = [
    tf.__internal__.distribute.combinations.tpu_strategy,
]

parameter_server_strategies_single_worker = [
    tf.__internal__.distribute.combinations.parameter_server_strategy_1worker_2ps_cpu,  # noqa: E501
    tf.__internal__.distribute.combinations.parameter_server_strategy_1worker_2ps_1gpu,  # noqa: E501
]

parameter_server_strategies_multi_worker = [
    tf.__internal__.distribute.combinations.parameter_server_strategy_3worker_2ps_cpu,  # noqa: E501
    tf.__internal__.distribute.combinations.parameter_server_strategy_3worker_2ps_1gpu,  # noqa: E501
]

all_strategies = strategies_minus_tpu + tpu_strategies