1# Copyright 2018 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"""This module customizes `test_combinations` for Tensorflow. 16 17Additionally it provides `generate()`, `combine()` and `times()` with Tensorflow 18customizations as a default. 19""" 20 21import functools 22 23from tensorflow.python import tf2 24from tensorflow.python.eager import context 25from tensorflow.python.framework import ops 26from tensorflow.python.framework import test_combinations 27from tensorflow.python.util.tf_export import tf_export 28 29 30class EagerGraphCombination(test_combinations.TestCombination): 31 """Run the test in Graph or Eager mode. 32 33 The optional `mode` parameter controls the test's execution mode. Its 34 accepted values are "graph" or "eager" literals. 35 """ 36 37 def context_managers(self, kwargs): 38 mode = kwargs.pop("mode", None) 39 if mode is None: 40 return [] 41 elif mode == "eager": 42 return [context.eager_mode()] 43 elif mode == "graph": 44 return [ops.Graph().as_default(), context.graph_mode()] 45 else: 46 raise ValueError( 47 "Argument 'mode' must be either 'eager' or 'graph'. " 48 f"Received: {mode}.") 49 50 def parameter_modifiers(self): 51 return [test_combinations.OptionalParameter("mode")] 52 53 54class TFVersionCombination(test_combinations.TestCombination): 55 """Control the execution of the test in TF1.x and TF2. 56 57 If TF2 is enabled then a test with TF1 test is going to be skipped and vice 58 versa. 59 60 Test targets continuously run in TF2 thanks to the tensorflow.v2 TAP target. 61 A test can be run in TF2 with bazel by passing --test_env=TF2_BEHAVIOR=1. 62 """ 63 64 def should_execute_combination(self, kwargs): 65 tf_api_version = kwargs.pop("tf_api_version", None) 66 if tf_api_version == 1 and tf2.enabled(): 67 return (False, "Skipping a TF1.x test when TF2 is enabled.") 68 elif tf_api_version == 2 and not tf2.enabled(): 69 return (False, "Skipping a TF2 test when TF2 is not enabled.") 70 return (True, None) 71 72 def parameter_modifiers(self): 73 return [test_combinations.OptionalParameter("tf_api_version")] 74 75 76generate = functools.partial( 77 test_combinations.generate, 78 test_combinations=(EagerGraphCombination(), TFVersionCombination())) 79combine = test_combinations.combine 80times = test_combinations.times 81NamedObject = test_combinations.NamedObject 82 83tf_export("__internal__.test.combinations.generate", v1=[])(generate) 84