def create_function_unit_test( function_name: str, unit_test_spec: tf_test_utils.UnitTestSpec) -> tf.function: """Creates a tf_function_unit_test from the provided UnitTestSpec.""" function = getattr(tf.math, function_name) signature = unit_test_spec.input_signature if tf_utils.is_complex(signature): function, signature = tf_utils.rewrite_complex_signature( function, signature) if function_name == "top_k": function = _wrap_top_k(function) wrapped_function = lambda *args: function(*args, **unit_test_spec.kwargs) if FLAGS.dynamic_dims: signature = tf_utils.apply_function(signature, tf_utils.make_dims_dynamic) return tf_test_utils.tf_function_unit_test( input_signature=signature, input_generator=unit_test_spec.input_generator, input_args=unit_test_spec.input_args, name=unit_test_spec.unit_test_name, rtol=1e-5, atol=1e-5)(wrapped_function)
def __init__(self): super().__init__() self.m = utils.get_model_with_default_params( FLAGS.model, MODE_ENUM_TO_MODE[FLAGS.mode]) call = lambda *args: self.m(*args, training=False) input_signature = [ tf.TensorSpec(tensor.shape) for tensor in self.m.inputs ] self.call = tf_test_utils.tf_function_unit_test( input_signature=input_signature, name="call", atol=1e-5)(call)
def __init__(self): super().__init__() self.m = initialize_model() input_shape = list([BATCH_SIZE] + self.m.inputs[0].shape[1:]) # Some models accept dynamic image dimensions by default, so we use # IMAGE_DIM as a stand-in. for i, dim in enumerate(input_shape): if dim is None: input_shape[i] = IMAGE_DIM # Specify input shape with a static batch size. # TODO(b/142948097): Add support for dynamic shapes in SPIR-V lowering. self.call = tf_test_utils.tf_function_unit_test( input_signature=[tf.TensorSpec(input_shape)], name="call", rtol=1e-5, atol=1e-5)(lambda x: self.m(x, training=False))
def create_layer_unit_test( model: tf.keras.Model, unit_test_spec: tf_test_utils.UnitTestSpec) -> tf.function: """Wrap the model's __call__ function in a tf.function for testing.""" static_signature = unit_test_spec.input_signature dynamic_signature = static_signature if FLAGS.dynamic_dims: dynamic_signature = tf_utils.apply_function(dynamic_signature, tf_utils.make_dims_dynamic) if len(static_signature) > 1: static_signature = [static_signature] dynamic_signature = [dynamic_signature] call = lambda *args: model(*args, training=FLAGS.training) return tf_test_utils.tf_function_unit_test( input_signature=dynamic_signature, static_signature=static_signature, input_generator=unit_test_spec.input_generator, input_args=unit_test_spec.input_args, name=unit_test_spec.unit_test_name)(call)