1# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Strategy combinations for combinations.combine()."""
16
17from tensorflow.python.distribute import strategy_combinations
18
19
20multidevice_strategies = [
21    strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
22    strategy_combinations.mirrored_strategy_with_two_gpus,
23    strategy_combinations.tpu_strategy,
24]
25
26multiworker_strategies = [
27    strategy_combinations.multi_worker_mirrored_2x1_cpu,
28    strategy_combinations.multi_worker_mirrored_2x1_gpu,
29    strategy_combinations.multi_worker_mirrored_2x2_gpu
30]
31
32strategies_minus_default_minus_tpu = [
33    strategy_combinations.one_device_strategy,
34    strategy_combinations.one_device_strategy_gpu,
35    strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
36    strategy_combinations.mirrored_strategy_with_two_gpus,
37    strategy_combinations.central_storage_strategy_with_gpu_and_cpu
38]
39
40strategies_minus_tpu = [
41    strategy_combinations.default_strategy,
42    strategy_combinations.one_device_strategy,
43    strategy_combinations.one_device_strategy_gpu,
44    strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
45    strategy_combinations.mirrored_strategy_with_two_gpus,
46    strategy_combinations.central_storage_strategy_with_gpu_and_cpu
47]
48
49multi_worker_mirrored_strategies = [
50    strategy_combinations.multi_worker_mirrored_2x1_cpu,
51    strategy_combinations.multi_worker_mirrored_2x1_gpu,
52    strategy_combinations.multi_worker_mirrored_2x2_gpu,
53]
54
55tpu_strategies = [
56    strategy_combinations.tpu_strategy,
57]
58
59all_strategies = strategies_minus_tpu + tpu_strategies
60