def benchmarkCudnnLSTMTraining(self):
    test_configs = self._GetTestConfig()
    for config_name, config in test_configs.items():
      config = test_configs[config_name]
      num_layers = config["num_layers"]
      num_units = config["num_units"]
      batch_size = config["batch_size"]
      seq_length = config["seq_length"]

      with ops.Graph().as_default(), ops.device("/gpu:0"):
        model = cudnn_rnn_ops.CudnnLSTM(num_layers, num_units, num_units)
        params_size_t = model.params_size()
        input_data = variables.Variable(
            array_ops.ones([seq_length, batch_size, num_units]))
        input_h = variables.Variable(
            array_ops.ones([num_layers, batch_size, num_units]))
        input_c = variables.Variable(
            array_ops.ones([num_layers, batch_size, num_units]))
        params = variables.Variable(
            array_ops.ones([params_size_t]), validate_shape=False)
        output, output_h, output_c = model(
            is_training=True,
            input_data=input_data,
            input_h=input_h,
            input_c=input_c,
            params=params)
        all_grads = gradients_impl.gradients(
            [output, output_h, output_c],
            [params, input_data, input_h, input_c])
        training_op = control_flow_ops.group(*all_grads)
        self._BenchmarkOp(training_op, "cudnn_lstm %s %s" %
                          (config_name, self._GetConfigDesc(config)))

  def testRegisterBlocks(self):
    with ops.Graph().as_default():
      random_seed.set_random_seed(200)
      lc = layer_collection.LayerCollection()
      lc.register_fully_connected(
          array_ops.constant(1), array_ops.constant(2), array_ops.constant(3))
      lc.register_fully_connected(
          array_ops.constant(1),
          array_ops.constant(2),
          array_ops.constant(3),
          approx=layer_collection.APPROX_DIAGONAL_NAME)
      lc.register_conv2d(
          array_ops.constant(4), [1, 1, 1, 1], 'SAME',
          array_ops.ones((1, 1, 1, 1)), array_ops.constant(3))
      lc.register_conv2d(
          array_ops.constant(4), [1, 1, 1, 1],
          'SAME',
          array_ops.ones((1, 1, 1, 1)),
          array_ops.constant(3),
          approx=layer_collection.APPROX_DIAGONAL_NAME)
      lc.register_generic(
          array_ops.constant(5), 16, approx=layer_collection.APPROX_FULL_NAME)
      lc.register_generic(
          array_ops.constant(6),
          16,
          approx=layer_collection.APPROX_DIAGONAL_NAME)

      self.assertEqual(6, len(lc.get_blocks()))

  def testRejectionDataListInput(self):
    batch_size = 20
    val_input_batch = [
        array_ops.zeros([2, 3, 4]), array_ops.ones([2, 4]), array_ops.ones(2) *
        3
    ]
    lbl_input_batch = array_ops.ones([], dtype=dtypes.int32)
    probs = np.array([0, 1, 0, 0, 0])
    val_list, lbls = sampling_ops.stratified_sample(
        val_input_batch,
        lbl_input_batch,
        probs,
        batch_size,
        init_probs=[0, 1, 0, 0, 0])

    # Check output shapes.
    self.assertTrue(isinstance(val_list, list))
    self.assertEqual(len(val_list), len(val_input_batch))
    self.assertTrue(isinstance(lbls, ops.Tensor))

    with self.test_session() as sess:
      coord = coordinator.Coordinator()
      threads = queue_runner_impl.start_queue_runners(coord=coord)

      out = sess.run(val_list + [lbls])

      coord.request_stop()
      coord.join(threads)

    # Check output shapes.
    self.assertEqual(len(out), len(val_input_batch) + 1)

def _lower_triangular_mask(shape):
  """Creates a lower-triangular boolean mask over the last 2 dimensions."""
  row_index = math_ops.cumsum(
      array_ops.ones(shape=shape, dtype=dtypes.int32), axis=-2)
  col_index = math_ops.cumsum(
      array_ops.ones(shape=shape, dtype=dtypes.int32), axis=-1)
  return math_ops.greater_equal(row_index, col_index)

  def testAcceptsTensor(self):
    tensor = array_ops.ones([10, 10])
    result = math_ops.scalar_mul(3, tensor)
    expected = array_ops.ones([10, 10]) * 3

    with self.test_session(use_gpu=True):
      self.assertAllEqual(expected.eval(), result.eval())

  def _variance(self):
    # We need to put the tf.where inside the outer tf.where to ensure we never
    # hit a NaN in the gradient.
    denom = array_ops.where(math_ops.greater(self.df, 2.),
                            self.df - 2.,
                            array_ops.ones_like(self.df))
    # Abs(scale) superfluous.
    var = (array_ops.ones(self.batch_shape_tensor(), dtype=self.dtype) *
           math_ops.square(self.scale) * self.df / denom)
    # When 1 < df <= 2, variance is infinite.
    inf = np.array(np.inf, dtype=self.dtype.as_numpy_dtype())
    result_where_defined = array_ops.where(
        self.df > array_ops.fill(self.batch_shape_tensor(), 2.),
        var,
        array_ops.fill(self.batch_shape_tensor(), inf, name="inf"))

    if self.allow_nan_stats:
      nan = np.array(np.nan, dtype=self.dtype.as_numpy_dtype())
      return array_ops.where(
          math_ops.greater(
              self.df,
              array_ops.ones(self.batch_shape_tensor(), dtype=self.dtype)),
          result_where_defined,
          array_ops.fill(self.batch_shape_tensor(), nan, name="nan"))
    else:
      return control_flow_ops.with_dependencies(
          [
              check_ops.assert_less(
                  array_ops.ones([], dtype=self.dtype),
                  self.df,
                  message="variance not defined for components of df <= 1"),
          ],
          result_where_defined)

  def testClusterSpecPropagationThreeServers2Graphs(self):
    """Boots 3 servers, creates 2 sessions, ensures appropriate operations.

    We create 2 clusterspecs:
     1. server2 as the master, server1 as a worker
     2. server2 as the master, server3 as a worker

    We ensure that variables on the workers are independent.
    """
    server1 = server_lib.Server.create_local_server()
    server2 = server_lib.Server.create_local_server()
    server3 = server_lib.Server.create_local_server()
    cluster_def1 = cluster_pb2.ClusterDef()
    job1 = cluster_def1.job.add()
    job1.name = 'worker1'
    job1.tasks[0] = server2.target[len('grpc://'):]
    job1.tasks[1] = server1.target[len('grpc://'):]

    cluster_def2 = cluster_pb2.ClusterDef()
    job2 = cluster_def2.job.add()
    job2.name = 'worker2'
    job2.tasks[0] = server2.target[len('grpc://'):]
    job2.tasks[1] = server3.target[len('grpc://'):]

    config1 = config_pb2.ConfigProto(cluster_def=cluster_def1)
    config2 = config_pb2.ConfigProto(cluster_def=cluster_def2)

    with ops.Graph().as_default() as g1:
      with ops.device('/job:worker1/task:1'):
        var1 = variables.Variable(array_ops.zeros([2]), name='var1')
        update_op1 = state_ops.assign_add(
            var1, array_ops.ones([2]), name='var1_assign_add')
        init1 = variables.global_variables_initializer()

    with ops.Graph().as_default() as g2:
      with ops.device('/job:worker2/task:1'):
        var2 = variables.Variable(array_ops.zeros([2]), name='var2')
        update_op2 = state_ops.assign_add(
            var2, array_ops.ones([2]), name='var2_assign_add')
        init2 = variables.global_variables_initializer()

    sess1 = session.Session(server2.target, graph=g1, config=config1)
    sess2 = session.Session(server2.target, graph=g2, config=config2)

    init1.run(session=sess1)
    init2.run(session=sess2)

    expected_zeros = np.zeros([2])
    expected_ones = np.ones([2])

    self.assertAllEqual(expected_zeros, sess1.run(var1))
    self.assertAllEqual(expected_zeros, sess2.run(var2))

    self.assertAllEqual(expected_ones, sess1.run(update_op1))
    self.assertAllEqual(expected_ones, sess1.run(var1))
    self.assertAllEqual(expected_zeros, sess2.run(var2))
    self.assertAllEqual(expected_ones, sess2.run(update_op2))
    self.assertAllEqual(expected_ones + expected_ones, sess1.run(update_op1))
    self.assertAllEqual(expected_ones, sess2.run(var2))
    self.assertAllEqual(expected_ones + expected_ones, sess1.run(var1))

 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 testAdamSparse(self):
    with ops.device('/cpu:0'):
      # Create 2-D embedding for 3 objects on CPU because sparse/sliced updates
      # are not implemented on TPU.
      embedding_matrix = resource_variable_ops.ResourceVariable(
          array_ops.ones([3, 2]))

    with self.test_scope():
      with backprop.GradientTape() as tape:
        embedding = embedding_ops.embedding_lookup(embedding_matrix, [1])
        y = math_ops.reduce_sum(embedding)
      dy_dx = tape.gradient(y, embedding_matrix)
      self.assertIsInstance(dy_dx, ops.IndexedSlices)
      optimizer = adam.AdamOptimizer(0.1)
      # The gradient application operations will run on CPU because optimizer
      # updates are always collocated with the variable.
      optimizer.apply_gradients([(dy_dx, embedding_matrix)])

      # This assign_add will run on CPU because when an input to an
      # operation is a resource, this operation is placed on the resource's
      # device by the eager runtime.
      embedding_matrix.assign_add(array_ops.ones([3, 2]))

    self.assertAllClose([[2.0, 2.0],
                         [1.9, 1.9],
                         [2.0, 2.0]], embedding_matrix.numpy())

 def _testOneSimpleInference(self, rnn_mode, num_layers, num_units, input_size,
                             batch_size, seq_length, dir_count, expected,
                             tolerance):
   model = self._CreateModel(rnn_mode, num_layers, num_units, input_size)
   has_input_c = (rnn_mode == "lstm")
   params_size_t = model.params_size()
   input_data = array_ops.ones([seq_length, batch_size, input_size])
   input_h = array_ops.ones([num_layers * dir_count, batch_size, num_units])
   params = variables.Variable(
       array_ops.ones([params_size_t]), validate_shape=False)
   if has_input_c:
     input_c = array_ops.ones([num_layers * dir_count, batch_size, num_units])
     output, output_h, output_c = model(
         input_data=input_data,
         input_h=input_h,
         input_c=input_c,
         params=params,
         is_training=False)
   else:
     output, output_h = model(
         input_data=input_data,
         input_h=input_h,
         params=params,
         is_training=False)
   output_sum = math_ops.reduce_sum(output)
   output_h_sum = math_ops.reduce_sum(output_h)
   total_sum = output_sum + output_h_sum
   if has_input_c:
     output_c_sum = math_ops.reduce_sum(output_c)
     total_sum += output_c_sum
   with self.test_session(use_gpu=True) as sess:
     sess.run(variables.global_variables_initializer())
     total_sum_v = sess.run([total_sum])
     self.assertAllClose(
         total_sum_v[0], expected, atol=tolerance, rtol=tolerance)

  def testCovariance(self):
    with self.test_session():
      vex = ds.VectorExponentialDiag(
          loc=array_ops.ones([2, 3], dtype=dtypes.float32))
      self.assertAllClose(
          np.diag(np.ones([3], dtype=np.float32)),
          vex.covariance().eval())

      vex = ds.VectorExponentialDiag(
          loc=array_ops.ones([3], dtype=dtypes.float32),
          scale_identity_multiplier=[3., 2.])
      self.assertAllEqual([2], vex.batch_shape)
      self.assertAllEqual([3], vex.event_shape)
      self.assertAllClose(
          np.array([[[3., 0, 0],
                     [0, 3, 0],
                     [0, 0, 3]],
                    [[2, 0, 0],
                     [0, 2, 0],
                     [0, 0, 2]]])**2.,
          vex.covariance().eval())

      vex = ds.VectorExponentialDiag(
          loc=array_ops.ones([3], dtype=dtypes.float32),
          scale_diag=[[3., 2, 1], [4, 5, 6]])
      self.assertAllEqual([2], vex.batch_shape)
      self.assertAllEqual([3], vex.event_shape)
      self.assertAllClose(
          np.array([[[3., 0, 0],
                     [0, 2, 0],
                     [0, 0, 1]],
                    [[4, 0, 0],
                     [0, 5, 0],
                     [0, 0, 6]]])**2.,
          vex.covariance().eval())

 def _mode(self):
     mode = (self.a - 1.0) / (self.a_b_sum - 2.0)
     if self.allow_nan_stats:
         nan = np.array(np.nan, dtype=self.dtype.as_numpy_dtype())
         return math_ops.select(
             math_ops.logical_and(math_ops.greater(self.a, 1.0), math_ops.greater(self.b, 1.0)),
             mode,
             array_ops.fill(self.batch_shape(), nan, name="nan"),
         )
     else:
         return control_flow_ops.with_dependencies(
             [
                 check_ops.assert_less(
                     array_ops.ones((), dtype=self.dtype),
                     self.a,
                     message="Mode not defined for components of a <= 1.",
                 ),
                 check_ops.assert_less(
                     array_ops.ones((), dtype=self.dtype),
                     self.b,
                     message="Mode not defined for components of b <= 1.",
                 ),
             ],
             mode,
         )

  def testAcceptsTensor(self):
    tensor = array_ops.ones([10, 10])
    result = math_ops.scalar_mul(3, tensor)
    expected = array_ops.ones([10, 10]) * 3

    with test_util.device(use_gpu=True):
      self.assertAllEqual(self.evaluate(expected), self.evaluate(result))

 def testShape(self):
   # Fully known shape.
   rnd = random_ops.random_gamma([150], 2.0)
   self.assertEqual([150], rnd.get_shape().as_list())
   rnd = random_ops.random_gamma([150], 2.0, beta=[3.0, 4.0])
   self.assertEqual([150, 2], rnd.get_shape().as_list())
   rnd = random_ops.random_gamma([150], array_ops.ones([1, 2, 3]))
   self.assertEqual([150, 1, 2, 3], rnd.get_shape().as_list())
   rnd = random_ops.random_gamma([20, 30], array_ops.ones([1, 2, 3]))
   self.assertEqual([20, 30, 1, 2, 3], rnd.get_shape().as_list())
   rnd = random_ops.random_gamma(
       [123], array_ops.placeholder(
           dtypes.float32, shape=(2,)))
   self.assertEqual([123, 2], rnd.get_shape().as_list())
   # Partially known shape.
   rnd = random_ops.random_gamma(
       array_ops.placeholder(
           dtypes.int32, shape=(1,)), array_ops.ones([7, 3]))
   self.assertEqual([None, 7, 3], rnd.get_shape().as_list())
   rnd = random_ops.random_gamma(
       array_ops.placeholder(
           dtypes.int32, shape=(3,)), array_ops.ones([9, 6]))
   self.assertEqual([None, None, None, 9, 6], rnd.get_shape().as_list())
   # Unknown shape.
   rnd = random_ops.random_gamma(
       array_ops.placeholder(dtypes.int32),
       array_ops.placeholder(dtypes.float32))
   self.assertIs(None, rnd.get_shape().ndims)
   rnd = random_ops.random_gamma([50], array_ops.placeholder(dtypes.float32))
   self.assertIs(None, rnd.get_shape().ndims)

 def testDtype(self):
   with self.test_session():
     d = array_ops.fill([2, 3], 12., name="fill")
     self.assertEqual(d.get_shape(), [2, 3])
     # Test default type for both constant size and dynamic size
     z = array_ops.ones([2, 3])
     self.assertEqual(z.dtype, dtypes_lib.float32)
     self.assertEqual([2, 3], z.get_shape())
     self.assertAllEqual(z.eval(), np.ones([2, 3]))
     z = array_ops.ones(array_ops.shape(d))
     self.assertEqual(z.dtype, dtypes_lib.float32)
     self.assertEqual([2, 3], z.get_shape())
     self.assertAllEqual(z.eval(), np.ones([2, 3]))
     # Test explicit type control
     for dtype in (dtypes_lib.float32, dtypes_lib.float64, dtypes_lib.int32,
                   dtypes_lib.uint8, dtypes_lib.int16, dtypes_lib.int8,
                   dtypes_lib.complex64, dtypes_lib.complex128,
                   dtypes_lib.int64, dtypes_lib.bool):
       z = array_ops.ones([2, 3], dtype=dtype)
       self.assertEqual(z.dtype, dtype)
       self.assertEqual([2, 3], z.get_shape())
       self.assertAllEqual(z.eval(), np.ones([2, 3]))
       z = array_ops.ones(array_ops.shape(d), dtype=dtype)
       self.assertEqual(z.dtype, dtype)
       self.assertEqual([2, 3], z.get_shape())
       self.assertAllEqual(z.eval(), np.ones([2, 3]))

  def benchmarkMatrixBandPartOp(self):
    for shape_ in self.shapes:
      for limits in (-1, -1), (-1, 0), (0, -1), (2, 2):
        with ops.Graph().as_default(), \
            session.Session() as sess, \
            ops.device("/cpu:0"):
          matrix = variables.Variable(array_ops.ones(shape_))
          band = array_ops.matrix_band_part(matrix, limits[0], limits[1])
          variables.global_variables_initializer().run()
          self.run_op_benchmark(
              sess,
              control_flow_ops.group(band),
              min_iters=10,
              name="matrix_band_part_cpu_{shape}_{limits}".format(
                  shape=shape_, limits=limits))

        if test_lib.is_gpu_available(True):
          with ops.Graph().as_default(), \
              session.Session() as sess, \
              ops.device("/gpu:0"):
            matrix = variables.Variable(array_ops.ones(shape_))
            band = array_ops.matrix_band_part(matrix, limits[0], limits[1])
            variables.global_variables_initializer().run()
            self.run_op_benchmark(
                sess,
                control_flow_ops.group(band),
                min_iters=10,
                name="matrix_band_part_gpu_{shape}_{limits}".format(
                    shape=shape_, limits=limits))

 def test_noise_decreasing(self):
   for dtype in [dtypes.float32, dtypes.float64]:
     with variable_scope.variable_scope(dtype.name):
       random_model = RandomStateSpaceModel(
           state_dimension=5, state_noise_dimension=4,
           configuration=state_space_model.StateSpaceModelConfiguration(
               dtype=dtype, num_features=1))
       random_model.initialize_graph()
       original_covariance = array_ops.diag(
           array_ops.ones(shape=[5], dtype=dtype))
       _, new_covariance, _ = random_model._exogenous_noise_decreasing(
           current_times=[[1]],
           exogenous_values=constant_op.constant([[-2.]], dtype=dtype),
           state=[
               -array_ops.ones(shape=[1, 5], dtype=dtype),
               original_covariance[None], [0]
           ])
       with self.cached_session() as session:
         variables.global_variables_initializer().run()
         evaled_new_covariance, evaled_original_covariance = session.run(
             [new_covariance[0], original_covariance])
         new_variances = numpy.diag(evaled_new_covariance)
         original_variances = numpy.diag(evaled_original_covariance)
         for i in range(5):
           self.assertLess(new_variances[i], original_variances[i])

  def testServerDefChanged(self):
    """Update server def, and run ops on new cluster."""
    context.set_server_def(
        server_def=get_server_def(
            ALT_JOB_NAME,
            local_server_port=0,
            remote_server_addresses=[
                self._cached_server1_target, self._cached_server2_target
            ],
            task_index=0))

    with ops.device("job:%s/replica:0/task:1/device:CPU:0" % ALT_JOB_NAME):
      x1 = array_ops.ones([2, 2])
    y = math_ops.matmul(x1, x1)
    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())

    # Set the server def back to JOB_NAME
    context.set_server_def(
        server_def=get_server_def(
            JOB_NAME,
            local_server_port=0,
            remote_server_addresses=[
                self._cached_server1_target, self._cached_server2_target
            ],
            task_index=0))

    with ops.device("job:%s/replica:0/task:1/device:CPU:0" % JOB_NAME):
      x1 = array_ops.ones([2, 2])
    y = math_ops.matmul(x1, x1)
    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())

  def test_nested_network_inside_network(self):
    inner_inputs = {
        'x1': keras.Input(shape=(1,)),
        'x2': keras.Input(shape=(1,))
    }
    inner_outputs = {
        'x1+x2':
            keras.layers.Add()([inner_inputs['x1'], inner_inputs['x2']]),
        'x1*x2':
            keras.layers.Multiply()([inner_inputs['x1'], inner_inputs['x2']])
    }
    inner_network = keras.engine.network.Network(inner_inputs, inner_outputs)

    inputs = [keras.Input(shape=(1,)), keras.Input(shape=(1,))]
    middle = inner_network({'x1': inputs[0], 'x2': inputs[1]})
    outputs = keras.layers.Add()([middle['x1+x2'], middle['x1*x2']])
    network = keras.engine.network.Network(inputs, outputs)

    network = keras.engine.network.Network.from_config(network.get_config())

    # Computes: `(x1+x2) + (x1*x2)`
    result_tensor = network(
        [array_ops.ones((1, 1), 'float32'),
         array_ops.ones((1, 1), 'float32')])
    result = self.evaluate(result_tensor)
    self.assertAllEqual(result, [[3.]])

    output_shape = network.compute_output_shape([(None, 1), (None, 1)])
    self.assertListEqual(output_shape.as_list(), [None, 1])

 def testEagerSingleOutputFloat32(self):
   with test_util.device(use_gpu=True):
     a = array_ops.ones((3, 3), dtype=dtypes.float32)
     x = array_ops.ones((3, 1), dtype=dtypes.float32)
     output = script_ops.eager_py_func(matmul, inp=[a, x], Tout=dtypes.float32)
     ret = self.evaluate(output)
     self.assertAllClose(ret, [[3.0], [3.0], [3.0]])