code: # 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.
  # ==============================================================================
  """Tests for TensorFlow 2.0 layer behavior."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import collections
  import itertools as it
  import os
  import sys
  import traceback
  from absl.testing import parameterized
  import numpy as np
  
  from tensorflow.python import keras
  from tensorflow.python.eager import context
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.framework import test_util
  from tensorflow.python.keras import backend
  from tensorflow.python.keras import keras_parameterized
  from tensorflow.python.keras import testing_utils
  from tensorflow.python.keras.engine import base_layer
  from tensorflow.python.keras.mixed_precision.experimental import policy
  from tensorflow.python.keras.optimizer_v2 import rmsprop
  from tensorflow.python.keras.utils import tf_utils
  from tensorflow.python.layers import core as legacy_core
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import state_ops
  from tensorflow.python.ops import summary_ops_v2
  from tensorflow.python.ops import tensor_array_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import gfile
  from tensorflow.python.platform import test
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.summary import summary_iterator
  from tensorflow.python.util import nest
  
  
  class DynamicLayer(base_layer.Layer):
  
    def __init__(self, dynamic=False, **kwargs):
      super(DynamicLayer, self).__init__(dynamic=dynamic, **kwargs)
  
    def call(self, inputs):
      samples = tensor_array_ops.TensorArray(
          dtype=dtypes.float32, size=array_ops.shape(inputs)[0])
      for idx, sample in enumerate(inputs):
        samples = samples.write(idx, math_ops.square(sample))
      return samples.stack()
  
    def compute_output_shape(self, input_shape):
      return input_shape
  
  
  class InvalidLayer(base_layer.Layer):
  
    def call(self, inputs):
      raise ValueError('You did something wrong!')
  
  
  class BaseLayerTest(keras_parameterized.TestCase):
  
    @keras_parameterized.run_with_all_model_types
    def test_dynamic_layer(self):
      model = testing_utils.get_model_from_layers([DynamicLayer(dynamic=True)],
                                                  input_shape=(3,))
      self.assertEqual(model.dynamic, True)
      model.compile(rmsprop.RMSprop(0.001), loss='mse')
      self.assertEqual(model.run_eagerly, True)
      model.train_on_batch(np.random.random((2, 3)), np.random.random((2, 3)))
  
    @keras_parameterized.run_with_all_model_types
    def test_dynamic_layer_error(self):
      with self.assertRaisesRegexp(TypeError,
                                   'attempting to use Python control flow'):
        model = testing_utils.get_model_from_layers([DynamicLayer()],
                                                    input_shape=(3,))
        model.compile(rmsprop.RMSprop(0.001), loss='mse')
        model.train_on_batch(np.random.random((2, 3)), np.random.random((2, 3)))
  
    @keras_parameterized.run_with_all_model_types
    def test_dynamic_layer_error_running_in_graph_mode(self):
      with context.graph_mode():
        model = testing_utils.get_model_from_layers([DynamicLayer(dynamic=True)],
                                                    input_shape=(3,))
        self.assertEqual(model.dynamic, True)
        # But then you cannot run the model since you're in a graph scope.
        with self.assertRaisesRegexp(
            ValueError, 'You must enable eager execution'):
          model.compile(rmsprop.RMSprop(0.001), loss='mse')
  
    def test_manual_compute_output_shape(self):
      class BuildCounter(keras.layers.Layer):
  
        def __init__(self, *args, **kwargs):  # pylint: disable=redefined-outer-name
          super(BuildCounter, self).__init__(*args, **kwargs)
          self.build_counter = 0
  
        def build(self, input_shape):
          self.build_counter += 1
  
        def call(self, inputs):
          return inputs
  
      with context.eager_mode():
        layer = BuildCounter(dtype=dtypes.float64)
        output_shape = layer.compute_output_shape((None, 10))
        self.assertEqual(layer.build_counter, 1)
        self.assertEqual(output_shape.as_list(), [None, 10])
        output_signature = layer.compute_output_signature(
            tensor_spec.TensorSpec(dtype=dtypes.float64, shape=[None, 10]))
        self.assertEqual(layer.build_counter, 1)
        self.assertEqual(output_signature.dtype, dtypes.float64)
        self.assertEqual(output_signature.shape.as_list(), [None, 10])
        layer(np.ones((5, 10)))
        self.assertEqual(layer.build_counter, 1)
  
    def test_dynamic_layer_with_deferred_sequential_model(self):
      model = keras.Sequential(
          [DynamicLayer(dynamic=True),
           keras.layers.Dense(3)])
      self.assertEqual(model.dynamic, True)
      model.compile(rmsprop.RMSprop(0.001), loss='mse')
      self.assertEqual(model.run_eagerly, True)
      model.train_on_batch(np.random.random((2, 3)), np.random.random((2, 3)))
  
    def test_nested_dynamic_layers_in_eager_mode(self):
      inputs = keras.Input((3,))
      outputs = DynamicLayer(dynamic=True)(inputs)
      inner_model = keras.Model(inputs, outputs)
      self.assertEqual(inner_model.dynamic, True)
  
      inputs = keras.Input((3,))
      x = DynamicLayer(dynamic=True)(inputs)
      outputs = inner_model(x)
  
      model = keras.Model(inputs, outputs)
      self.assertEqual(model.dynamic, True)
      model.compile(rmsprop.RMSprop(0.001), loss='mse')
      self.assertEqual(model.run_eagerly, True)
      model.train_on_batch(np.random.random((2, 3)), np.random.random((2, 3)))
  
    def test_dynamic_subclassed_model_no_shape_inference(self):
  
      class MyModel(keras.Model):
  
        def __init__(self):
          super(MyModel, self).__init__(dynamic=True)
          self.layer1 = keras.layers.Dense(3)
          self.layer2 = keras.layers.Dense(3)
  
        def call(self, inputs):
          if math_ops.reduce_sum(inputs) > 0:
            return self.layer1(inputs)
          else:
            return self.layer2(inputs)
  
      model = MyModel()
      self.assertEqual(model.dynamic, True)
      model.compile(rmsprop.RMSprop(0.001), loss='mse')
      self.assertEqual(model.run_eagerly, True)
      model.train_on_batch(np.random.random((2, 3)), np.random.random((2, 3)))
      self.assertEqual(model.outputs, [None])
  
    def test_dynamic_subclassed_model_with_shape_inference(self):
  
      class MyModel(keras.Model):
  
        def __init__(self):
          super(MyModel, self).__init__(dynamic=True)
          self.layer1 = keras.layers.Dense(3)
          self.layer2 = keras.layers.Dense(3)
  
        def call(self, inputs):
          if math_ops.reduce_sum(inputs) > 0:
            return self.layer1(inputs)
          else:
            return self.layer2(inputs)
  
        def compute_output_shape(self, input_shape):
          return tensor_shape.TensorShape(
              tuple(input_shape[:-1].as_list()) + (3,))
  
      model = MyModel()
      self.assertEqual(model.dynamic, True)
      model.compile(rmsprop.RMSprop(0.001), loss='mse')
      model.train_on_batch(np.random.random((2, 3)), np.random.random((2, 3)))
      self.assertEqual(model.outputs[0].shape.as_list(), [None, 3])
  
    @keras_parameterized.run_all_keras_modes
    def test_add_loss_correctness(self):
  
      class MyLayer(keras.layers.Layer):
  
        def call(self, inputs, training=None):
          self.add_loss(math_ops.reduce_sum(inputs))
          return inputs
  
      inputs = keras.Input((3,))
      layer = MyLayer()
      outputs = layer(inputs)
      model = keras.Model(inputs, outputs)
      self.assertEqual(len(model.losses), 1)
      model.compile(
          'sgd',
          'mse',
          run_eagerly=testing_utils.should_run_eagerly(),
          experimental_run_tf_function=testing_utils.should_run_tf_function())
      loss = model.train_on_batch(np.ones((2, 3)), np.ones((2, 3)))
      self.assertEqual(loss, 2 * 3)
  
    @test_util.run_in_graph_and_eager_modes
    def test_invalid_forward_pass(self):
      inputs = keras.Input((3,))
      with self.assertRaisesRegexp(ValueError, 'You did something wrong!'):
        _ = InvalidLayer()(inputs)
  
    def test_no_legacy_model(self):
      inputs = keras.Input((1,))
      legacy_dense_0 = legacy_core.Dense(1, name='legacy_dense_0')
      legacy_dense_1 = legacy_core.Dense(1, name='legacy_dense_1')
  
      layer = legacy_dense_0(inputs)
      layer = keras.layers.Dense(1)(layer)
      layer = legacy_dense_1(layer)
  
      expected_regex = (r'The following are legacy tf\.layers\.Layers:\n  '
                        '{}\n  {}'.format(legacy_dense_0, legacy_dense_1))
  
      with self.assertRaisesRegexp(TypeError, expected_regex):
        _ = keras.models.Model(inputs=[inputs], outputs=[layer])
  
      model = keras.models.Model(inputs=[inputs], outputs=[inputs])
      with self.assertRaisesRegexp(TypeError, expected_regex):
        model._insert_layers([legacy_dense_0, legacy_dense_1])
  
    def test_no_legacy_sequential(self):
      layers = [
          keras.layers.Dense(1),
          legacy_core.Dense(1, name='legacy_dense_0')
      ]
  
      expected_regex = r'legacy tf\.layers\.Layers:\n  {}'.format(layers[1])
      with self.assertRaisesRegexp(TypeError, expected_regex):
        _ = keras.models.Sequential(layers)
  
      with self.assertRaisesRegexp(TypeError, expected_regex):
        _ = keras.models.Sequential([keras.layers.Input(shape=(4,))] + layers)
  
      model = keras.models.Sequential()
      with self.assertRaisesRegexp(TypeError, expected_regex):
        for l in layers:
          model.add(l)
  
    @keras_parameterized.run_with_all_model_types
    @test_util.run_in_graph_and_eager_modes
    def test_build_with_numpy_data(self):
      model_layers = [
          keras.layers.Dense(3, activation='relu', kernel_initializer='ones'),
          keras.layers.Dense(1, activation='sigmoid', kernel_initializer='ones')
      ]
      model = testing_utils.get_model_from_layers(model_layers, input_shape=(4,))
      model(np.zeros((2, 4), dtype='float32'))
      self.assertTrue(model.built)
  
    @test_util.run_in_graph_and_eager_modes
    def test_default_add_weight(self):
  
      class TestLayer(keras.layers.Layer):
  
        def __init__(self):
          super(TestLayer, self).__init__()
          self.default_weight = self.add_weight()
          self.weight_without_name = self.add_weight(shape=(3, 4))
          self.regularized_weight_without_name = self.add_weight(
              shape=(3, 4), regularizer='l2')
  
      layer = TestLayer()
      self.assertEqual(layer.default_weight.shape.as_list(), [])
      self.assertEqual(layer.weight_without_name.shape.as_list(), [3, 4])
      self.assertEqual(layer.default_weight.dtype.name, 'float32')
      self.assertEqual(layer.weight_without_name.dtype.name, 'float32')
      self.assertEqual(len(layer.losses), 1)
      if not context.executing_eagerly():
        # Cannot access tensor.name in eager execution.
        self.assertTrue('Variable_2/Regularizer' in layer.losses[0].name)
  
    @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
    def test_learning_phase_freezing_for_layers(self):
      class LearningPhaseLayer(keras.layers.Layer):
  
        def call(self, inputs):
          return keras.backend.in_train_phase(
              lambda: array_ops.ones_like(inputs),
              lambda: array_ops.zeros_like(inputs))
  
      def get_learning_phase_value():
        model = keras.models.Sequential([LearningPhaseLayer(input_shape=(1,))])
        model._run_eagerly = testing_utils.should_run_eagerly()
        model._experimental_run_tf_function = (
            testing_utils.should_run_tf_function())
        return np.sum(model(np.ones((1, 1))))
  
      self.assertEqual(get_learning_phase_value(), 0)
  
      # Test scope.
      with keras.backend.learning_phase_scope(1):
        self.assertEqual(get_learning_phase_value(), 1)
  
      # The effects of the scope end after exiting it.
      self.assertEqual(get_learning_phase_value(), 0)
  
      # Test setting.
      keras.backend.set_learning_phase(1)
      self.assertEqual(get_learning_phase_value(), 1)
      keras.backend.set_learning_phase(0)
      self.assertEqual(get_learning_phase_value(), 0)
  
    @keras_parameterized.run_all_keras_modes
    def test_learning_phase_freezing_for_layers_in_predict(self):
      if not (testing_utils.should_run_eagerly() or
              testing_utils.should_run_tf_function()):
        self.skipTest('Predict fails to override the outer learning phase in'
                      'the FuncGraph path.')
  
      class LearningPhaseLayer(keras.layers.Layer):
  
        def call(self, inputs):
          return keras.backend.in_train_phase(
              lambda: array_ops.ones_like(inputs),
              lambda: array_ops.zeros_like(inputs))
  
      def get_learning_phase_value():
        model = keras.models.Sequential([LearningPhaseLayer(input_shape=(1,))])
        model._run_eagerly = testing_utils.should_run_eagerly()
        model._experimental_run_tf_function = (
            testing_utils.should_run_tf_function())
        return np.sum(model.predict(np.ones((1, 1))))
  
      self.assertEqual(get_learning_phase_value(), 0)
  
      # Test scope.
      with keras.backend.learning_phase_scope(1):
        self.assertEqual(get_learning_phase_value(), 0)
  
      # The effects of the scope end after exiting it.
      self.assertEqual(get_learning_phase_value(), 0)
  
      # Test setting.
      keras.backend.set_learning_phase(1)
      self.assertEqual(get_learning_phase_value(), 0)
      keras.backend.set_learning_phase(0)
      self.assertEqual(get_learning_phase_value(), 0)
  
    # Cannot be enabled with `run_eagerly=True`, see b/123904578
    @test_util.run_all_in_graph_and_eager_modes
    def test_layer_can_return_variable(self):
  
      class ComputeSum(keras.layers.Layer):
  
        def __init__(self):
          super(ComputeSum, self).__init__()
          self.total = variables.Variable(
              initial_value=array_ops.zeros((1, 1)), trainable=False)
          if not context.executing_eagerly():
            keras.backend.get_session().run(self.total.initializer)
  
        def call(self, inputs):
          self.total.assign_add(inputs)
          return self.total
  
      inputs = keras.Input(shape=(1,))
      model = keras.Model(inputs, ComputeSum()(inputs))
      model.predict(np.ones((1, 1)))
  
    def _get_layer_with_training_arg(self):
  
      class TrainingLayer(keras.layers.Layer):
        """A layer with a `training` argument in a defuned `call`."""
  
        @def_function.function
        def call(self, inputs, training=None):
          if training is None:
            training = keras.backend.learning_phase()
          return tf_utils.smart_cond(training,
                                     lambda: array_ops.ones_like(inputs),
                                     lambda: array_ops.zeros_like(inputs))
  
      return TrainingLayer()
  
    @keras_parameterized.run_with_all_model_types
    # b/124459427: can't test with `run_eagerly=True` for now.
    @test_util.run_in_graph_and_eager_modes
    def test_training_arg_in_defun(self):
      layer = self._get_layer_with_training_arg()
      model = testing_utils.get_model_from_layers([layer], input_shape=(1,))
      model.compile(rmsprop.RMSprop(0.),
                    loss='mae')
      history = model.fit(np.zeros((1, 1)), np.zeros((1, 1)))
      self.assertEqual(history.history['loss'][0], 1.)
      loss = model.evaluate(np.zeros((1, 1)), np.zeros((1, 1)))
      self.assertEqual(loss, 0.)
  
      # Test that the argument injection performed in `call` is not active
      # when the argument is passed explicitly.
      layer = self._get_layer_with_training_arg()
      inputs = keras.Input(shape=(1,))
      # Pass `training` by name
      outputs = layer(inputs, training=False)
      model = keras.Model(inputs, outputs)
      model.compile(rmsprop.RMSprop(0.),
                    loss='mae')
      history = model.fit(np.zeros((1, 1)), np.zeros((1, 1)))
      self.assertEqual(history.history['loss'][0], 0.)
  
    @keras_parameterized.run_with_all_model_types
    @keras_parameterized.run_all_keras_modes
    def test_raw_variable_assignment(self):
  
      class RawVariableLayer(keras.layers.Layer):
  
        def __init__(self, **kwargs):
          super(RawVariableLayer, self).__init__(**kwargs)
          # Test variables in nested structure.
          self.var_list = [variables.Variable(1.), {'a': variables.Variable(2.)}]
  
        def call(self, inputs):
          return inputs * self.var_list[0] * self.var_list[1]['a']
  
      model = testing_utils.get_model_from_layers([RawVariableLayer()],
                                                  input_shape=(10,))
      model.compile(
          'sgd',
          'mse',
          run_eagerly=testing_utils.should_run_eagerly(),
          experimental_run_tf_function=testing_utils.should_run_tf_function())
      x, y = np.ones((10, 10)), np.ones((10, 10))
      # Checks that variables get initialized.
      model.fit(x, y, batch_size=2, epochs=2)
  
    @test_util.run_in_graph_and_eager_modes
    def test_layer_names(self):
      inputs = keras.layers.Input(shape=[2])
      add1 = inputs + inputs
      add2 = keras.layers.Add()([inputs, inputs])
      add3 = inputs + inputs
      add4 = keras.layers.Add()([inputs, inputs])
      model = keras.models.Model(
          inputs=[inputs], outputs=[add1, add2, add3, add4])
      self.assertEqual(
          [l.name for l in model.layers],
          ['input_1', 'tf_op_layer_add', 'add', 'tf_op_layer_add_2', 'add_1'])
  
    def test_add_trainable_weight_on_frozen_layer(self):
  
      class TestLayer(keras.layers.Layer):
  
        def build(self, input_shape):
          self.w = self.add_weight(shape=(), trainable=True)
  
        def call(self, inputs):
          return self.w * inputs
  
      layer = TestLayer()
      layer.trainable = False
      layer.build(None)
      layer.trainable = True
      self.assertListEqual(layer.trainable_weights, [layer.w])
  
    @keras_parameterized.run_with_all_model_types
    @keras_parameterized.run_all_keras_modes
    def test_passing_initial_weights_values(self):
      kernel_value = np.random.random((10, 2))
      layer_with_weights = keras.layers.Dense(
          2, use_bias=False, weights=[kernel_value])
  
      model = testing_utils.get_model_from_layers([layer_with_weights],
                                                  input_shape=(10,))
      model.compile(
          'sgd',
          'mse',
          run_eagerly=testing_utils.should_run_eagerly(),
          experimental_run_tf_function=testing_utils.should_run_tf_function())
      inputs = np.random.random((3, 10))
      out = model.predict(inputs)
      self.assertAllClose(model.layers[-1].get_weights()[0], kernel_value)
      self.assertAllClose(out, np.dot(inputs, kernel_value))
  
    @test_util.run_in_graph_and_eager_modes
    def test_set_weights_and_get_weights(self):
      layer = keras.layers.Dense(2)
      layer.build((None, 10))
      kernel = np.random.random((10, 2))
      bias = np.random.random((2,))
      layer.set_weights([kernel, bias])
      weights = layer.get_weights()
      self.assertEqual(len(weights), 2)
      self.assertAllClose(weights[0], kernel)
      self.assertAllClose(weights[1], bias)
      with self.assertRaisesRegexp(
          ValueError, 'but the layer was expecting 2 weights'):
        layer.set_weights([1, 2, 3])
      with self.assertRaisesRegexp(
          ValueError, 'not compatible with provided weight shape'):
        layer.set_weights([kernel.T, bias])
  
    def test_get_config_error(self):
  
      class MyLayer(keras.layers.Layer):
  
        def __init__(self, my_kwarg='default', **kwargs):
          super(MyLayer, self).__init__(**kwargs)
          self.my_kwarg = my_kwarg
  
      # `__init__` includes kwargs but `get_config` is not overridden, so
      # an error should be thrown:
      with self.assertRaises(NotImplementedError):
        MyLayer('custom').get_config()
  
      class MyLayerNew(keras.layers.Layer):
  
        def __init__(self, my_kwarg='default', **kwargs):
          super(MyLayerNew, self).__init__(**kwargs)
          self.my_kwarg = my_kwarg
  
        def get_config(self):
          config = super(MyLayerNew, self).get_config()
          config['my_kwarg'] = self.my_kwarg
          return config
  
      # Test to make sure that error is not raised if the method call is
      # from an overridden `get_config`:
      self.assertEqual(MyLayerNew('custom').get_config()['my_kwarg'], 'custom')
  
      class MyLayerNew2(keras.layers.Layer):
  
        def __init__(self, name='MyLayerName', dtype=None, **kwargs):  # pylint:disable=redefined-outer-name
          super(MyLayerNew2, self).__init__(name=name, dtype=dtype, **kwargs)
  
      # Check that if the kwargs in `__init__` are base layer constructor
      # arguments, no error is thrown:
      self.assertEqual(MyLayerNew2(name='New').get_config()['name'], 'New')
  
  
  class SymbolicSupportTest(test.TestCase):
  
    def test_using_symbolic_tensors_with_tf_ops(self):
      # Single-input.
      x = keras.Input((3,))
      y = math_ops.square(x)
      self.assertEqual(y.graph, keras.backend.get_graph())
  
      # Multi-inputs.
      x1, x2 = keras.Input((3,)), keras.Input((3,))
      y = array_ops.concat([x1, x2], axis=1)
      self.assertEqual(y.graph, keras.backend.get_graph())
  
      # Mixing Keras symbolic tensors and graph tensors from the same graph works.
      with keras.backend.get_graph().as_default():
        x1 = keras.Input((3,))
      x2 = keras.Input((3,))
      y = math_ops.matmul(x1, x2)
      self.assertEqual(y.graph, keras.backend.get_graph())
  
      # Creating same op type (matmul) multiple times in the Keras graph works.
      x1 = keras.Input((3,))
      x2 = keras.Input((3,))
      y = math_ops.matmul(x1, x2)
      self.assertEqual(y.graph, keras.backend.get_graph())
  
    def test_mixing_eager_and_graph_tensors(self):
      with ops.Graph().as_default():
        x1 = array_ops.ones((3, 3))
      x2 = array_ops.ones((3, 3))
      self.assertIsInstance(x2, ops.EagerTensor)
      with self.assertRaisesRegexp(TypeError, 'Graph tensors'):
        math_ops.matmul(x1, x2)
  
    def test_mixing_numpy_arrays_and_graph_tensors(self):
      with ops.Graph().as_default():
        x1 = array_ops.ones((3, 3))
      x2 = np.ones((3, 3), dtype='float32')
      with self.assertRaisesRegexp(TypeError, 'Graph tensors'):
        math_ops.matmul(x1, x2)
  
    @test_util.run_in_graph_and_eager_modes
    def test_mixing_keras_symbolic_tensors_and_eager_tensors(self):
      x1 = keras.Input((3,))
      x2 = array_ops.ones((3, 3))
      y = math_ops.matmul(x1, x2)
      self.assertEqual(y.graph, keras.backend.get_graph())
      fn = keras.backend.function(inputs=[x1], outputs=[y])
      x_val = np.random.random((3, 3))
      y_val = np.ones((3, 3))
      self.assertAllClose(fn([x_val])[0],
                          np.matmul(x_val, y_val),
                          atol=1e-5)
  
    @test_util.run_in_graph_and_eager_modes
    def test_mixing_keras_symbolic_tensors_and_numpy_arrays(self):
      x1 = keras.Input((3,))
      x2 = np.ones((3, 3), dtype='float32')
      y = math_ops.matmul(x1, x2)
      self.assertEqual(y.graph, keras.backend.get_graph())
      fn = keras.backend.function(inputs=[x1], outputs=[y])
      x_val = np.random.random((3, 3))
      y_val = np.ones((3, 3))
      self.assertAllClose(fn([x_val])[0],
                          np.matmul(x_val, y_val),
                          atol=1e-5)
  
    @test_util.run_in_graph_and_eager_modes
    def test_reraising_exception(self):
      # When layer is not dynamic, we have some pattern matching during exception
      # handling to detect when the user is trying to use python control flow.
      # When an exception is thrown but the pattern doesn't match, we want to
      # preserve the originating stack trace. An early implementation of this
      # logic lost the stack trace. We test the correct behavior here.
  
      class TypeErrorLayer(base_layer.Layer):
  
        def call(self, inputs):
          def easily_identifiable_name():
            raise TypeError('Non-matching TypeError message.')
          easily_identifiable_name()
  
      inputs = keras.Input((3,))
  
      try:
        _ = TypeErrorLayer()(inputs)
      except TypeError as e:
        if hasattr(e, 'ag_error_metadata'):
          self.assertIn('easily_identifiable_name', str(e))
          # See ErrorMetadataBase in autograph/pyct/errors.py
          # Topmost frame corresponds to `call` itself.
          function_name = e.ag_error_metadata.translated_stack[-2].function_name
        else:
          tb = traceback.extract_tb(sys.exc_info()[2])
          last_entry = tb[-1]
          function_name = last_entry[2]
        self.assertEqual(function_name, 'easily_identifiable_name')
  
    @test_util.run_in_graph_and_eager_modes
    def test_summaries_in_tf_function(self):
      if not context.executing_eagerly():
        return
  
      class MyLayer(keras.layers.Layer):
  
        def call(self, inputs):
          summary_ops_v2.scalar('mean', math_ops.reduce_mean(inputs))
          return inputs
  
      tmp_dir = self.get_temp_dir()
      writer = summary_ops_v2.create_file_writer_v2(tmp_dir)
      with writer.as_default(), summary_ops_v2.always_record_summaries():
        my_layer = MyLayer()
        x = array_ops.ones((10, 10))
  
        def my_fn(x):
          return my_layer(x)
  
        _ = my_fn(x)
  
      event_file = gfile.Glob(os.path.join(tmp_dir, 'events*'))
      self.assertLen(event_file, 1)
      event_file = event_file[0]
      tags = set()
      for e in summary_iterator.summary_iterator(event_file):
        for val in e.summary.value:
          tags.add(val.tag)
      self.assertEqual(set(['my_layer/mean']), tags)
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class NestedTrackingTest(test.TestCase):
  
    def test_nested_layer_variable_tracking(self):
      # Test that variables from nested sublayers are
      # being tracked by subclassed layers.
  
      class MyLayer(keras.layers.Layer):
  
        def __init__(self):
          super(MyLayer, self).__init__()
          self.dense1 = keras.layers.Dense(1)
          self.dense2 = keras.layers.BatchNormalization()
  
        def build(self, input_shape):
          self.v1 = self.add_weight('v1', shape=input_shape[1:].as_list())
          self.v2 = variables.Variable(
              name='v2',
              initial_value=np.zeros(input_shape[1:].as_list(), dtype='float32'),
              trainable=False)
  
        def call(self, inputs):
          x = self.dense1(inputs) + self.dense2(inputs)
          return x + self.v1 + self.v2
  
      layer = MyLayer()
      inputs = keras.Input((1,))
      _ = layer(inputs)
  
      self.assertEqual(len(layer.weights), 8)
      self.assertEqual(len(layer.trainable_weights), 5)
      self.assertEqual(len(layer.non_trainable_weights), 3)
  
      layer.dense1.trainable = False
      self.assertEqual(len(layer.weights), 8)
      self.assertEqual(len(layer.trainable_weights), 3)
      self.assertEqual(len(layer.non_trainable_weights), 5)
  
      layer.trainable = False
      self.assertEqual(len(layer.weights), 8)
      self.assertEqual(len(layer.trainable_weights), 0)
      self.assertEqual(len(layer.non_trainable_weights), 8)
      self.assertEqual(
          set([layer.dense1, layer.dense2, layer.v1, layer.v2]),
          set([obj for unused_name, obj in layer._checkpoint_dependencies]))
  
    def test_nested_layer_updates_losses_tracking(self):
      # Test that updates and losses from nested sublayers are
      # being tracked by subclassed layers.
  
      class UpdateAndLossLayer(keras.layers.Layer):
  
        def build(self, _):
          self.v1 = self.add_weight('v1', shape=())
  
        def call(self, inputs):
          self.add_loss(math_ops.reduce_sum(inputs))
          self.add_update(state_ops.assign_add(self.v1, 1))
          return inputs + 1
  
      class MyLayer(keras.layers.Layer):
  
        def build(self, _):
          self.v1 = self.add_weight('v1', shape=())
  
        def __init__(self):
          super(MyLayer, self).__init__()
          self.ul1 = UpdateAndLossLayer()
          self.ul2 = UpdateAndLossLayer()
  
        def call(self, inputs):
          self.add_loss(math_ops.reduce_sum(inputs))
          self.add_update(state_ops.assign_add(self.v1, 1))
          x = self.ul1(inputs)
          return self.ul2(x)
  
      layer = MyLayer()
  
      if context.executing_eagerly():
        inputs = array_ops.ones((3, 1))
        _ = layer(inputs)
        self.assertEqual(len(layer.losses), 3)
        self.assertLen(layer.get_losses_for(None), 3)
      else:
        inputs = keras.Input((1,))
        _ = layer(inputs)
        self.assertEqual(len(layer.losses), 3)
        self.assertEqual(len(layer.updates), 3)
        self.assertLen(layer.get_losses_for(None), 3)
  
    def test_attribute_reassignment(self):
      l = keras.layers.Layer()
      l.a = keras.layers.Layer()
      l.a = []
      l.a = variables.Variable(1.)
      l.a = keras.layers.Layer()
      last_assignment = keras.layers.Layer()
      l.a = last_assignment
      l.b = variables.Variable(1.)
      del l.b
      l.c = keras.layers.Layer()
      del l.c
      l.d = last_assignment
      del l.d
      self.assertEqual([last_assignment], l._layers)
      self.assertEqual([], l.trainable_weights)
      self.assertEqual([], l.non_trainable_weights)
      self.assertEqual([], l.weights)
      del l.a
      self.assertEqual([], l._layers)
  
    def test_assign_op_not_tracked_as_variable(self):
  
      class LayerWithAssignAttr(keras.layers.Layer):
  
        def build(self, input_shape):
          self.v = variables.Variable(1.)
          self.v_assign = self.v.assign_add(2.)
  
      layer = LayerWithAssignAttr()
      layer.build((10, 10))
  
      self.assertEqual([layer.v], layer.variables)
  
    def test_layer_class_not_tracked_as_sublayer(self):
      # See https://github.com/tensorflow/tensorflow/issues/27431 for details.
  
      class LayerWithClassAttribute(keras.layers.Layer):
  
        def __init__(self):
          super(LayerWithClassAttribute, self).__init__()
          self.layer_fn = keras.layers.Dense
  
      layer = LayerWithClassAttribute()
      self.assertEmpty(layer.variables)
      self.assertEmpty(layer.submodules)
  
    def test_layer_call_fn_args(self):
  
      class NonDefunLayer(keras.layers.Layer):
  
        def call(self, inputs, a, mask, b=None, training=None):
          return inputs
  
      class DefunLayer(keras.layers.Layer):
  
        @def_function.function
        def call(self, x, mask, a, training=None, b=None):
          return x
  
      nondefun_layer = NonDefunLayer()
      self.assertEqual(nondefun_layer._call_fn_args,
                       ['inputs', 'a', 'mask', 'b', 'training'])
      defun_layer = DefunLayer()
      self.assertEqual(defun_layer._call_fn_args,
                       ['x', 'mask', 'a', 'training', 'b'])
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class NameScopingTest(keras_parameterized.TestCase):
  
    def test_name_scope_layer(self):
      x = keras.backend.placeholder(shape=(10, 10))
      layer = keras.layers.Dense(10, name='MyName')
      layer(x)
      self.assertEqual(layer.bias.name, 'MyName/bias:0')
      self.assertEqual(layer.kernel.name, 'MyName/kernel:0')
  
    def test_name_scope_sublayer(self):
  
      class NameScopeTracker(keras.layers.Layer):
  
        def call(self, inputs):
          self.active_name_scope = ops.get_name_scope()
          return inputs
  
      x = keras.backend.placeholder(shape=(10, 10))
      sublayer = NameScopeTracker(name='Sublayer')
      layer = keras.layers.Dense(10, activation=sublayer, name='MyName2')
      layer(x)
      self.assertEqual(layer.bias.name, 'MyName2/bias:0')
      self.assertEqual(layer.kernel.name, 'MyName2/kernel:0')
      self.assertEqual(sublayer.active_name_scope, 'MyName2/Sublayer')
  
    def test_name_scope_tf_tensor(self):
      x = ops.convert_to_tensor(np.ones((10, 10)))
      layer = keras.layers.Dense(
          10, activation=keras.layers.ReLU(name='MyAct'), name='MyName3')
      layer(x)
      self.assertEqual(layer.bias.name, 'MyName3/bias:0')
      self.assertEqual(layer.kernel.name, 'MyName3/kernel:0')
  
  
  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
  class AutographControlFlowTest(keras_parameterized.TestCase):
  
    def test_disabling_in_context_is_matched(self):
  
      test_obj = self
  
      class MyLayer(keras.layers.Layer):
  
        def call(self, inputs, training=None):
          with test_obj.assertRaisesRegex(TypeError, 'Tensor.*as.*bool'):
            if constant_op.constant(False):
              return inputs * 1.
          return inputs * 0.
  
      @def_function.function(autograph=False)
      def test_fn():
        return MyLayer()(constant_op.constant([[1., 2., 3.]]))
  
      test_fn()
  
    def test_if_training_pattern_output(self):
  
      class MyLayer(keras.layers.Layer):
  
        def call(self, inputs, training=None):
          if training:
            return inputs * 1.
          return inputs * 0.
  
      inputs = keras.Input((3,))
      outputs = MyLayer()(inputs)
      model = keras.Model(inputs, outputs)
      model.compile(
          'sgd',
          'mse',
          run_eagerly=testing_utils.should_run_eagerly(),
          experimental_run_tf_function=testing_utils.should_run_tf_function())
      train_loss = model.train_on_batch(np.ones((2, 3)), np.ones((2, 3)))
      self.assertEqual(train_loss, 0.)
      test_loss = model.test_on_batch(np.ones((2, 3)), np.ones((2, 3)))
      self.assertEqual(test_loss, 1.)
  
    def test_if_training_pattern_loss(self):
  
      class MyLayer(keras.layers.Layer):
  
        def call(self, inputs, training=None):
          if training:
            loss = math_ops.reduce_sum(inputs)
          else:
            loss = 0.
          self.add_loss(loss)
          return inputs
  
      inputs = keras.Input((3,))
      outputs = MyLayer()(inputs)
      model = keras.Model(inputs, outputs)
      model.compile(
          'sgd',
          'mse',
          run_eagerly=testing_utils.should_run_eagerly(),
          experimental_run_tf_function=testing_utils.should_run_tf_function())
      train_loss = model.train_on_batch(np.ones((2, 3)), np.ones((2, 3)))
      self.assertEqual(train_loss, 2 * 3)
      test_loss = model.test_on_batch(np.ones((2, 3)), np.ones((2, 3)))
      self.assertEqual(test_loss, 0)
  
    def test_if_training_pattern_metric(self):
  
      class MyLayer(keras.layers.Layer):
  
        def call(self, inputs, training=None):
          if training:
            metric = math_ops.reduce_sum(inputs)
          else:
            metric = 0.
          self.add_metric(metric, name='my_metric', aggregation='mean')
          return inputs
  
      inputs = keras.Input((3,))
      outputs = MyLayer()(inputs)
      model = keras.Model(inputs, outputs)
      model.compile(
          'sgd',
          'mse',
          run_eagerly=testing_utils.should_run_eagerly(),
          experimental_run_tf_function=testing_utils.should_run_tf_function())
      _, train_metric = model.train_on_batch(np.ones((2, 3)),
                                             np.ones((2, 3)))
      self.assertEqual(train_metric, 2 * 3)
      _, test_metric = model.test_on_batch(np.ones((2, 3)),
                                           np.ones((2, 3)))
      self.assertEqual(test_metric, 0)
  
    def test_if_training_pattern_update(self):
  
      class MyLayer(keras.layers.Layer):
  
        def build(self, input_shape):
          self.counter = self.add_weight(
              shape=(), trainable=False, initializer='zeros')
  
        def call(self, inputs, training=None):
          if training:
            increment = 1.
          else:
            increment = 0.
          self.counter.assign_add(increment)
          return inputs
  
      inputs = keras.Input((3,))
      layer = MyLayer()
      outputs = layer(inputs)
      model = keras.Model(inputs, outputs)
      model.compile(
          'sgd',
          'mse',
          run_eagerly=testing_utils.should_run_eagerly(),
          experimental_run_tf_function=testing_utils.should_run_tf_function())
      model.train_on_batch(np.ones((2, 3)), np.ones((2, 3)))
      self.assertEqual(keras.backend.get_value(layer.counter), 1.)
  
    def test_conditional_updates_in_call(self):
  
      class MyLayer(keras.layers.Layer):
  
        def __init__(self):
          super(MyLayer,
                self).__init__(dynamic=testing_utils.should_run_eagerly())
  
        def build(self, input_shape):
          self.counter = self.add_weight(
              shape=(), trainable=False, initializer='zeros')
  
        def call(self, inputs, training=None):
          if training:
            z = math_ops.reduce_sum(inputs)
            self.add_update(lambda: self.counter.assign_add(z))
          return inputs
  
        def compute_output_shape(self, input_shape):
          return input_shape
  
      if testing_utils.should_run_eagerly():
        inputs = keras.Input((3,))
        layer = MyLayer()
        outputs = layer(inputs)
        model = keras.Model(inputs, outputs)
        model.compile(
            'sgd',
            'mse',
            run_eagerly=testing_utils.should_run_eagerly(),
            experimental_run_tf_function=testing_utils.should_run_tf_function())
        model.train_on_batch(np.ones((2, 3)), np.ones((2, 3)))
        self.assertEqual(keras.backend.get_value(layer.counter), 6.)
      else:
        # TODO(fchollet): support the same workflow in graph mode.
        with self.assertRaisesRegexp(RuntimeError,
                                     '`add_update` in a control flow branch'):
          layer = MyLayer()
          layer(keras.Input((3,)))
          _ = layer.updates
  
    def test_conditional_losses_in_call(self):
  
      class MyLayer(keras.layers.Layer):
  
        def __init__(self):
          super(MyLayer,
                self).__init__(dynamic=testing_utils.should_run_eagerly())
  
        def call(self, inputs, training=None):
          if training:
            self.add_loss(math_ops.reduce_sum(inputs))
          return inputs
  
        def compute_output_shape(self, input_shape):
          return input_shape
  
      if testing_utils.should_run_eagerly():
        inputs = keras.Input((3,))
        layer = MyLayer()
        outputs = layer(inputs)
        model = keras.Model(inputs, outputs)
        model.compile(
            'sgd',
            'mse',
            run_eagerly=testing_utils.should_run_eagerly(),
            experimental_run_tf_function=testing_utils.should_run_tf_function())
        loss = model.train_on_batch(np.ones((2, 3)), np.ones((2, 3)))
        self.assertEqual(loss, 2 * 3)
      else:
        with self.assertRaisesRegexp(RuntimeError,
                                     '`add_loss` in a control flow branch'):
          layer = MyLayer()(keras.Input((3,)))
  
    def test_conditional_callable_losses(self):
      model = keras.Sequential([
          keras.layers.Dense(
              1, kernel_regularizer=keras.regularizers.l2(1e-4), input_shape=(1,))
      ])
      model._run_eagerly = testing_utils.should_run_eagerly()
      model._experimental_run_tf_function = testing_utils.should_run_tf_function()
  
      def assert_graph(t):
        if not context.executing_eagerly():
          self.assertEqual(t.graph, ops.get_default_graph())
  
      @def_function.function
      def get_losses(t):
        if t < 0:
          return math_ops.reduce_sum(model.losses) * t
        else:
          return math_ops.reduce_sum(model.losses)
  
      assert_graph(get_losses(constant_op.constant(2.)))
      assert_graph(get_losses(constant_op.constant(0.5)))
  
    def test_conditional_metrics_in_call(self):
  
      class MyLayer(keras.layers.Layer):
  
        def __init__(self):
          super(MyLayer,
                self).__init__(dynamic=testing_utils.should_run_eagerly())
  
        def call(self, inputs, training=None):
          if training:
            self.add_metric(math_ops.reduce_sum(inputs),
                            name='sum',
                            aggregation='mean')
          return inputs
  
        def compute_output_shape(self, input_shape):
          return input_shape
  
      if testing_utils.should_run_eagerly():
        inputs = keras.Input((3,))
        layer = MyLayer()
        outputs = layer(inputs)
        model = keras.Model(inputs, outputs)
        model.compile(
            'sgd',
            'mse',
            run_eagerly=testing_utils.should_run_eagerly(),
            experimental_run_tf_function=testing_utils.should_run_tf_function())
        history = model.fit(np.ones((2, 3)), np.ones((2, 3)))
        self.assertEqual(history.history['sum'][-1], 2 * 3)
      else:
        # TODO(fchollet): support the same workflow in graph mode.
        with self.assertRaisesRegexp(RuntimeError,
                                     '`add_metric` in a control flow branch'):
          layer = MyLayer()(keras.Input((3,)))
  
    def test_conditional_activity_regularizer_in_call(self):
  
      class TestModel(keras.Model):
  
        def __init__(self):
          super(TestModel, self).__init__(
              name='test_model', dynamic=testing_utils.should_run_eagerly())
          self.layer = keras.layers.Dense(2, activity_regularizer='l2')
  
        def call(self, x, training=None):
          if math_ops.greater(math_ops.reduce_sum(x), 0.0):
            return self.layer(x)
          else:
            return self.layer(x)
  
      model = TestModel()
      model.compile(
          loss='mse',
          optimizer='sgd',
          run_eagerly=testing_utils.should_run_eagerly(),
          experimental_run_tf_function=testing_utils.should_run_tf_function())
  
      x = np.ones(shape=(10, 1))
      y = np.ones(shape=(10, 2))
  
      if testing_utils.should_run_eagerly():
        model.fit(x, y, epochs=2, batch_size=5)
      else:
        with self.assertRaisesRegexp(
            RuntimeError, '`activity_regularizer` in a control flow branch'):
          model.fit(x, y, epochs=2, batch_size=5)
  
    def test_conditional_activity_regularizer_with_wrappers_in_call(self):
  
      class TestModel(keras.Model):
  
        def __init__(self):
          super(TestModel, self).__init__(
              name='test_model', dynamic=testing_utils.should_run_eagerly())
          self.layer = keras.layers.TimeDistributed(
              keras.layers.Dense(2, activity_regularizer='l2'),
              input_shape=(3, 4))
  
        def call(self, x, training=None):
          if math_ops.greater(math_ops.reduce_sum(x), 0.0):
            return self.layer(x)
          else:
            return self.layer(x)
  
      model = TestModel()
      model.compile(
          loss='mse',
          optimizer='sgd',
          run_eagerly=testing_utils.should_run_eagerly(),
          experimental_run_tf_function=testing_utils.should_run_tf_function())
  
      x = np.ones(shape=(10, 3, 4))
      y = np.ones(shape=(10, 3, 2))
  
      if testing_utils.should_run_eagerly():
        model.fit(x, y, epochs=2, batch_size=5)
      else:
        with self.assertRaisesRegexp(
            RuntimeError, '`activity_regularizer` in a control flow branch'):
          model.fit(x, y, epochs=2, batch_size=5)
  
  
  class AddLayer(keras.layers.Layer):
    """A layer which adds it's input to a variable.
  
    Useful for testing a layer with a variable
    """
  
    def build(self, _):
      self.v = self.add_weight('v', (), initializer='ones')
      self.built = True
  
    def call(self, inputs):
      return inputs + self.v
  
  
  class IdentityLayer(keras.layers.Layer):
    """A layer that returns it's input.
  
    Useful for testing a layer without a variable.
    """
  
    def call(self, inputs):
      return inputs
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class DTypeTest(keras_parameterized.TestCase):
  
    # This class only have tests relating to layer.dtype. Tests for dtype policies
    # are in mixed_precision/experimental/keras_test.py
  
    def _const(self, dtype):
      return array_ops.constant(1, dtype=dtype)
  
    @testing_utils.enable_v2_dtype_behavior
    def test_dtype_defaults_to_floatx(self):
      layer = AddLayer()
      self.assertEqual(layer.dtype, 'float32')
      layer(self._const('float64'))
      self.assertEqual(layer.dtype, 'float32')  # dtype should not change
  
      try:
        backend.set_floatx('float64')
        layer = AddLayer()
        self.assertEqual(layer.dtype, 'float64')
      finally:
        backend.set_floatx('float32')
  
    @testing_utils.enable_v2_dtype_behavior
    def test_passing_dtype_to_constructor(self):
      layer = IdentityLayer(dtype='float64')
      layer(self._const('float32'))
      self.assertEqual(layer.dtype, 'float64')
  
      layer = IdentityLayer(dtype='int32')
      layer(self._const('float32'))
      self.assertEqual(layer.dtype, 'int32')
  
      layer = IdentityLayer(dtype=dtypes.float64)
      layer(self._const('float32'))
      self.assertEqual(layer.dtype, 'float64')
  
    @testing_utils.enable_v2_dtype_behavior
    def input_cast_to_dtype(self):
      layer = AddLayer()
  
      # Input should be cast to layer.dtype, so output should also be layer.dtype
      self.assertEqual(layer(self._const('float64')).dtype, 'float32')
  
      layer = AddLayer(dtype='float64')
      self.assertEqual(layer(self._const('float32')).dtype, 'float64')
  
      # Test inputs are not casted if layer.dtype is not floating-point
      layer = IdentityLayer(dtype='int32')
      self.assertEqual(layer(self._const('float64')).dtype, 'float64')
  
      # Test inputs are not casted if the inputs are not floating-point
      layer = IdentityLayer(dtype='float32')
      self.assertEqual(layer(self._const('int32')).dtype, 'int32')
  
      # Test Numpy arrays are casted
      layer = IdentityLayer(dtype='float64')
      self.assertEqual(layer(np.array(1, dtype='float32')).dtype, 'float64')
  
      # Test Python floats are casted
      layer = IdentityLayer(dtype='float64')
      self.assertEqual(layer(1.).dtype, 'float64')
  
    @testing_utils.enable_v2_dtype_behavior
    def multiple_inputs_cast_to_dtype(self):
  
      class MultiIdentityLayer(keras.layers.Layer):
  
        def call(self, inputs):
          return [array_ops.identity(x) for x in inputs]
  
      # Testing layer with default dtype of float32
      layer = MultiIdentityLayer()
      x, y = layer([self._const('float16'), self._const('float32')])
      self.assertEqual(x.dtype, 'float32')
      self.assertEqual(y.dtype, 'float32')
  
      # Test passing dtype to the constructor
      layer = MultiIdentityLayer(dtype='float64')
      x, y = layer([self._const('float16'), self._const('float32')])
      self.assertEqual(x.dtype, 'float64')
      self.assertEqual(y.dtype, 'float64')
  
      # Test several non-floating point types
      layer = MultiIdentityLayer(dtype='float64')
      x, y, z, w = layer([self._const('float16'), self._const('bool'),
                          self._const('float64'), self._constant('complex64')])
      self.assertEqual(x.dtype, 'float64')
      self.assertEqual(y.dtype, 'bool')
      self.assertEqual(z.dtype, 'float64')
      self.assertEqual(w.dtype, 'complex64')
  
    @testing_utils.enable_v2_dtype_behavior
    def test_extra_args_and_kwargs_not_casted(self):
  
      class IdentityLayerWithArgs(keras.layers.Layer):
  
        def call(self, inputs, *args, **kwargs):
          return nest.flatten([inputs, args, kwargs])
  
      layer = IdentityLayerWithArgs(dtype='float64')
      x, y, z = layer(self._const('float16'), self._const('float16'),
                      kwarg=self._const('float16'))
      self.assertEqual(x.dtype, 'float64')
      self.assertEqual(y.dtype, 'float16')
      self.assertEqual(z.dtype, 'float16')
  
    @testing_utils.enable_v2_dtype_behavior
    def test_layer_without_autocast(self):
  
      class IdentityLayerWithoutAutocast(IdentityLayer):
  
        def __init__(self, *args, **kwargs):
          kwargs['experimental_autocast'] = False
          super(IdentityLayerWithoutAutocast, self).__init__(*args, **kwargs)
  
      layer = IdentityLayerWithoutAutocast(dtype='float64')
      self.assertEqual(layer(self._const('float32')).dtype, 'float32')
  
    @testing_utils.enable_v2_dtype_behavior
    def test_dtype_warnings(self):
      # Test a layer warns when it casts inputs.
      layer = IdentityLayer()
      with test.mock.patch.object(tf_logging, 'warn') as mock_warn:
        layer(self._const('float64'))
        self.assertRegexpMatches(
            str(mock_warn.call_args),
            ".*from dtype float64 to the layer's dtype of float32.*"
            "The layer has dtype float32 because.*")
  
      # Test a layer does not warn a second time
      with test.mock.patch.object(tf_logging, 'warn') as mock_warn:
        layer(self._const('float64'))
        mock_warn.assert_not_called()
  
      # Test a new layer can warn even if a different layer already warned
      layer = IdentityLayer()
      with test.mock.patch.object(tf_logging, 'warn') as mock_warn:
        layer(self._const('float64'))
        self.assertRegexpMatches(
            str(mock_warn.call_args),
            ".*from dtype float64 to the layer's dtype of float32.*"
            "The layer has dtype float32 because.*")
  
      # Test a layer does not warn if a dtype is passed
      layer = IdentityLayer(dtype='float32')
      with test.mock.patch.object(tf_logging, 'warn') as mock_warn:
        layer(self._const('float64'))
        mock_warn.assert_not_called()
  
      # Test a layer does not warn if a Policy is set:
      with policy.policy_scope('float32'):
        layer = IdentityLayer()
        with test.mock.patch.object(tf_logging, 'warn') as mock_warn:
          layer(self._const('float64'))
          mock_warn.assert_not_called()
  
    @testing_utils.enable_v2_dtype_behavior
    def test_compute_output_signature(self):
  
      class IdentityLayerWithOutputShape(IdentityLayer):
  
        def compute_output_shape(self, input_shape):
          return input_shape
  
      layer = IdentityLayerWithOutputShape(dtype='float64')
      output_signature = layer.compute_output_signature(
          tensor_spec.TensorSpec(shape=(), dtype='float32'))
      self.assertEqual(output_signature.shape, ())
      self.assertEqual(output_signature.dtype, 'float64')
  
    @testing_utils.enable_v2_dtype_behavior
    def test_passing_non_tensor(self):
      layer = IdentityLayer()
      x = object()
      y = layer(x)  # Layer should not cast 'x', as it's not a tensor
      self.assertIs(x, y)
  
    @testing_utils.disable_v2_dtype_behavior
    def test_v1_behavior(self):
      # Test dtype defaults to None and inferred from input
      layer = IdentityLayer()
      self.assertIsNone(layer.dtype)
      layer(self._const('float64'))
      self.assertEqual(layer.dtype, 'float64')
  
      # Test layer does not cast to dtype
      self.assertEqual(layer(self._const('float32')).dtype, 'float32')
  
  _LAYERS_TO_TEST = [
      (keras.layers.Dense, (1,), collections.OrderedDict(units=[1])),
      (keras.layers.Activation, (2, 2),
       collections.OrderedDict(activation=['relu'])),
      (keras.layers.Dropout, (16,), collections.OrderedDict(rate=[0.25])),
      (keras.layers.BatchNormalization, (8, 8, 3), collections.OrderedDict(
          axis=[3], center=[True, False], scale=[True, False])),
      (keras.layers.Conv1D, (8, 8), collections.OrderedDict(
          filters=[1], kernel_size=[1, 3], strides=[1, 2],
          padding=['valid', 'same'], use_bias=[True, False],
          kernel_regularizer=[None, 'l2'])),
      (keras.layers.Conv2D, (8, 8, 3), collections.OrderedDict(
          filters=[1], kernel_size=[1, 3], strides=[1, 2],
          padding=['valid', 'same'], use_bias=[True, False],
          kernel_regularizer=[None, 'l2'])),
      (keras.layers.LSTM, (8, 8), collections.OrderedDict(
          units=[1],
          activation=[None, 'relu'],
          kernel_regularizer=[None, 'l2'],
          dropout=[0, 0.5],
          stateful=[True, False],
          unroll=[True, False])),
  ]
  
  OUTPUT_TEST_CASES = []
  for layer_type, inp_shape, arg_dict in _LAYERS_TO_TEST:
    arg_combinations = [[(k, i) for i in v] for k, v in arg_dict.items()]  # pylint: disable=g-complex-comprehension
    for arguments in it.product(*arg_combinations):
      name = '_{}_{}'.format(layer_type.__name__,
                             '_'.join('{}_{}'.format(k, v) for k, v in arguments))
      OUTPUT_TEST_CASES.append(
          (name, layer_type, inp_shape, {k: v for k, v in arguments}))
  
  
  class OutputTypeTest(keras_parameterized.TestCase):
    """Test that layers and models produce the correct tensor types."""
  
    # In v1 graph there are only symbolic tensors.
    @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
    @parameterized.named_parameters(*OUTPUT_TEST_CASES)
    def test_layer_outputs(self, layer_to_test, input_shape, layer_kwargs):
      layer = layer_to_test(**layer_kwargs)
  
      input_data = np.ones(shape=(2,) + input_shape, dtype=np.float32)
      layer_result = layer(input_data)
  
      inp = keras.layers.Input(shape=input_shape, batch_size=2)
      model = keras.models.Model(inp, layer_to_test(**layer_kwargs)(inp))
      model_result = model(input_data)
  
      for x in [layer_result, model_result]:
        if not isinstance(x, ops.Tensor):
          raise ValueError('Tensor or EagerTensor expected, got type {}'
                           .format(type(x)))
  
        if isinstance(x, ops.EagerTensor) != context.executing_eagerly():
          expected_type = (ops.EagerTensor if context.executing_eagerly()
                           else ops.Tensor)
          raise ValueError('Expected type {}, got type {}'
                           .format(expected_type, type(x)))
  
  
  if __name__ == '__main__':
    ops.enable_eager_execution()
    test.main()
  
apis: numpy.ones tensorflow.python.eager.context.executing_eagerly tensorflow.python.keras.layers.Add tensorflow.python.framework.tensor_spec.TensorSpec tensorflow.python.keras.backend.get_value tensorflow.python.util.nest.flatten tensorflow.python.keras.testing_utils.get_model_from_layers tensorflow.python.keras.testing_utils.should_run_eagerly tensorflow.python.eager.def_function.function tensorflow.python.keras.backend.learning_phase tensorflow.python.framework.constant_op.constant tensorflow.python.keras.backend.function tensorflow.python.summary.summary_iterator.summary_iterator tensorflow.python.framework.ops.Graph tensorflow.python.keras.models.Sequential tensorflow.python.keras.backend.placeholder tensorflow.python.ops.math_ops.matmul tensorflow.python.eager.context.eager_mode tensorflow.python.keras.mixed_precision.experimental.policy.policy_scope tensorflow.python.framework.ops.get_name_scope tensorflow.python.ops.state_ops.assign_add tensorflow.python.keras.layers.ReLU tensorflow.python.keras.Model tensorflow.python.framework.ops.enable_eager_execution tensorflow.python.keras.Input tensorflow.python.ops.math_ops.square tensorflow.python.ops.summary_ops_v2.always_record_summaries numpy.zeros numpy.dot tensorflow.python.keras.backend.set_floatx tensorflow.python.ops.array_ops.identity tensorflow.python.ops.math_ops.reduce_sum tensorflow.python.keras.keras_parameterized.run_all_keras_modes tensorflow.python.ops.array_ops.shape tensorflow.python.ops.array_ops.zeros tensorflow.python.ops.summary_ops_v2.create_file_writer_v2 tensorflow.python.ops.math_ops.reduce_mean tensorflow.python.keras.backend.get_graph tensorflow.python.keras.layers.Layer tensorflow.python.ops.array_ops.constant tensorflow.python.eager.context.graph_mode tensorflow.python.keras.optimizer_v2.rmsprop.RMSprop tensorflow.python.keras.backend.learning_phase_scope tensorflow.python.keras.regularizers.l2 tensorflow.python.ops.variables.Variable tensorflow.python.ops.array_ops.ones numpy.matmul tensorflow.python.keras.layers.Dense tensorflow.python.layers.core.Dense tensorflow.python.framework.ops.get_default_graph tensorflow.python.ops.array_ops.zeros_like tensorflow.python.keras.models.Model numpy.random.random tensorflow.python.platform.test.main tensorflow.python.ops.array_ops.ones_like tensorflow.python.keras.layers.BatchNormalization tensorflow.python.ops.array_ops.concat tensorflow.python.keras.backend.get_session numpy.array tensorflow.python.keras.backend.set_learning_phase tensorflow.python.keras.testing_utils.should_run_tf_function tensorflow.python.keras.layers.Input tensorflow.python.platform.test.mock.patch.object