def testBatchFunctionOpWithCapturedInput(self):
    """Tests that batch_function op works with captured input."""
    with self.test_session() as sess:
      captured_inp0 = array_ops.placeholder_with_default(2, shape=[])
      captured_inp1 = array_ops.placeholder_with_default(1, shape=[])
      inp = array_ops.placeholder(dtype=dtypes.int32, shape=[1])

      @function.Defun(dtypes.int32)
      def computation(inp):
        return inp + captured_inp0 - captured_inp1

      result = gen_batch_ops.batch_function(
          num_batch_threads=1,
          max_batch_size=10,
          batch_timeout_micros=100000,  # 100ms
          allowed_batch_sizes=[3, 10],
          batching_queue="",
          f=computation,
          in_tensors=[inp],
          captured_tensors=computation.captured_inputs,
          Tout=[o.type for o in computation.definition.signature.output_arg])

      thread_results = []

      def worker():
        thread_results.extend(sess.run([result], feed_dict={inp: [1]}))

      worker_thread = threading.Thread(target=worker)
      worker_thread.start()
      main_results = sess.run([result], feed_dict={inp: [2]})
      worker_thread.join()
      self.assertEqual(thread_results[0], [2])
      self.assertEqual(main_results[0], [3])

  def test_bad_reshape_size(self):
    dims = 2
    new_batch_shape = [2, 3]
    old_batch_shape = [2]   # 2 != 2*3

    new_batch_shape_ph = (
        constant_op.constant(np.int32(new_batch_shape)) if self.is_static_shape
        else array_ops.placeholder_with_default(
            np.int32(new_batch_shape), shape=None))

    scale = np.ones(old_batch_shape + [dims], self.dtype)
    scale_ph = array_ops.placeholder_with_default(
        scale, shape=scale.shape if self.is_static_shape else None)
    mvn = mvn_lib.MultivariateNormalDiag(scale_diag=scale_ph)

    if self.is_static_shape:
      with self.assertRaisesRegexp(
          ValueError, (r"`batch_shape` size \(6\) must match "
                       r"`distribution\.batch_shape` size \(2\)")):
        batch_reshape_lib.BatchReshape(
            distribution=mvn,
            batch_shape=new_batch_shape_ph,
            validate_args=True)

    else:
      with self.test_session():
        with self.assertRaisesOpError(r"Shape sizes do not match."):
          batch_reshape_lib.BatchReshape(
              distribution=mvn,
              batch_shape=new_batch_shape_ph,
              validate_args=True).sample().eval()

  def testError(self,
                descr,
                mode,
                data,
                repeats,
                axis,
                exception=ValueError,
                error=None):
    # Make sure that this is also an error case for numpy.
    with self.assertRaises(exception):
      np.repeat(data, repeats, axis)

    if mode == 'constant':
      data = constant_op.constant(data)
      repeats = constant_op.constant(repeats)
    elif mode == 'dynamic':
      data = constant_op.constant(data)
      repeats = constant_op.constant(repeats)
      data = array_ops.placeholder_with_default(data, data.shape)
      repeats = array_ops.placeholder_with_default(repeats, repeats.shape)
    elif mode == 'unknown_shape':
      data = array_ops.placeholder_with_default(data, None)
      repeats = array_ops.placeholder_with_default(repeats, None)

    with self.assertRaisesRegexp(exception, error):
      ragged_util.repeat(data, repeats, axis)

  def test_non_positive_shape(self):
    dims = 2
    old_batch_shape = [4]
    if self.is_static_shape:
      # Unknown first dimension does not trigger size check. Note that
      # any dimension < 0 is treated statically as unknown.
      new_batch_shape = [-1, 0]
    else:
      new_batch_shape = [-2, -2]  # -2 * -2 = 4, same size as the old shape.

    new_batch_shape_ph = (
        constant_op.constant(np.int32(new_batch_shape)) if self.is_static_shape
        else array_ops.placeholder_with_default(
            np.int32(new_batch_shape), shape=None))

    scale = np.ones(old_batch_shape + [dims], self.dtype)
    scale_ph = array_ops.placeholder_with_default(
        scale, shape=scale.shape if self.is_static_shape else None)
    mvn = mvn_lib.MultivariateNormalDiag(scale_diag=scale_ph)

    if self.is_static_shape:
      with self.assertRaisesRegexp(ValueError, r".*must be >=-1.*"):
        batch_reshape_lib.BatchReshape(
            distribution=mvn,
            batch_shape=new_batch_shape_ph,
            validate_args=True)

    else:
      with self.test_session():
        with self.assertRaisesOpError(r".*must be >=-1.*"):
          batch_reshape_lib.BatchReshape(
              distribution=mvn,
              batch_shape=new_batch_shape_ph,
              validate_args=True).sample().eval()

  def test_non_vector_shape(self):
    dims = 2
    new_batch_shape = 2
    old_batch_shape = [2]

    new_batch_shape_ph = (
        constant_op.constant(np.int32(new_batch_shape)) if self.is_static_shape
        else array_ops.placeholder_with_default(
            np.int32(new_batch_shape), shape=None))

    scale = np.ones(old_batch_shape + [dims], self.dtype)
    scale_ph = array_ops.placeholder_with_default(
        scale, shape=scale.shape if self.is_static_shape else None)
    mvn = mvn_lib.MultivariateNormalDiag(scale_diag=scale_ph)

    if self.is_static_shape:
      with self.assertRaisesRegexp(ValueError, r".*must be a vector.*"):
        batch_reshape_lib.BatchReshape(
            distribution=mvn,
            batch_shape=new_batch_shape_ph,
            validate_args=True)

    else:
      with self.test_session():
        with self.assertRaisesOpError(r".*must be a vector.*"):
          batch_reshape_lib.BatchReshape(
              distribution=mvn,
              batch_shape=new_batch_shape_ph,
              validate_args=True).sample().eval()

  def test_broadcasting_explicitly_unsupported(self):
    old_batch_shape = [4]
    new_batch_shape = [1, 4, 1]
    rate_ = self.dtype([1, 10, 2, 20])

    rate = array_ops.placeholder_with_default(
        rate_,
        shape=old_batch_shape if self.is_static_shape else None)
    poisson_4 = poisson_lib.Poisson(rate)
    new_batch_shape_ph = (
        constant_op.constant(np.int32(new_batch_shape)) if self.is_static_shape
        else array_ops.placeholder_with_default(
            np.int32(new_batch_shape), shape=None))
    poisson_141_reshaped = batch_reshape_lib.BatchReshape(
        poisson_4, new_batch_shape_ph, validate_args=True)

    x_4 = self.dtype([2, 12, 3, 23])
    x_114 = self.dtype([2, 12, 3, 23]).reshape(1, 1, 4)

    if self.is_static_shape:
      with self.assertRaisesRegexp(NotImplementedError,
                                   "too few batch and event dims"):
        poisson_141_reshaped.log_prob(x_4)
      with self.assertRaisesRegexp(NotImplementedError,
                                   "unexpected batch and event shape"):
        poisson_141_reshaped.log_prob(x_114)
      return

    with self.assertRaisesOpError("too few batch and event dims"):
      with self.test_session():
        poisson_141_reshaped.log_prob(x_4).eval()

    with self.assertRaisesOpError("unexpected batch and event shape"):
      with self.test_session():
        poisson_141_reshaped.log_prob(x_114).eval()

  def testBatchDecoratedWithCapturedInput(self):
    """Tests that the batch_function decorator works."""
    if context.executing_eagerly():
      return
    with self.cached_session() as sess:
      captured_inp0 = array_ops.placeholder_with_default(2, shape=[])
      captured_inp1 = array_ops.placeholder_with_default(1, shape=[])

      @batch_ops.batch_function(1, 10, 100000)
      def computation(in_t):
        return in_t + captured_inp0 - captured_inp1

      inp = array_ops.placeholder(dtype=dtypes.int32, shape=[1])
      result = computation(inp)
      thread_results = []

      def worker():
        thread_results.extend(sess.run([result], feed_dict={inp: [1]}))

      worker_thread = threading.Thread(target=worker)
      worker_thread.start()
      main_results = sess.run([result], feed_dict={inp: [2]})
      worker_thread.join()
      self.assertEqual(thread_results[0], [2])
      self.assertEqual(main_results[0], [3])

  def testCDFWithDynamicEventShapeUnknownNdims(
      self, events, histograms, expected_cdf):
    """Test that dynamically-sized events with unknown shape work."""
    event_ph = array_ops.placeholder_with_default(events, shape=None)
    histograms_ph = array_ops.placeholder_with_default(histograms, shape=None)
    dist = categorical.Categorical(probs=histograms_ph)
    cdf_op = dist.cdf(event_ph)

    actual_cdf = self.evaluate(cdf_op)
    self.assertAllClose(actual_cdf, expected_cdf)

 def testRaggedTensorSplitsMismatchErrorAtRuntime(self):
   splits1 = array_ops.placeholder_with_default(
       constant_op.constant([0, 3, 3, 5], dtypes.int64), None)
   splits2 = array_ops.placeholder_with_default(
       constant_op.constant([0, 1, 3, 5], dtypes.int64), None)
   x = ragged_tensor.RaggedTensor.from_row_splits([3, 1, 4, 1, 5], splits1)
   y = ragged_tensor.RaggedTensor.from_row_splits([1, 2, 3, 4, 5], splits2)
   with self.assertRaisesRegexp(errors.InvalidArgumentError,
                                r'.*Inputs must have identical ragged splits'):
     self.evaluate(ragged_functional_ops.map_flat_values(math_ops.add, x, y))

 def setUp(self):
   super(BatchSequencesWithStatesTest, self).setUp()
   self.value_length = 4
   ind1 = np.array([
       [0, 0],
       [1, 0], [1, 3], [1, 4],
       [3, 2], [3, 3]])
   val1 = np.array([0, 10, 13, 14, 32, 33])
   shape1 = np.array([self.value_length, 6])
   sp_tensor1 = sparse_tensor.SparseTensor(
       array_ops.constant(ind1, dtypes.int64),
       array_ops.constant(val1, dtypes.int64),
       array_ops.placeholder_with_default(shape1, shape=[2]))
   ind2 = np.array([
       [0, 0, 1],
       [0, 1, 0],
       [0, 1, 2],
       [1, 0, 3],
       [1, 1, 0],
       [1, 1, 1],
       [1, 1, 2],
       [1, 2, 2]])
   val2 = np.array([1, 10, 12, 103, 150, 149, 150, 122])
   shape2 = np.array([self.value_length, 3, 4])
   sp_tensor2 = sparse_tensor.SparseTensor(
       array_ops.constant(ind2, dtypes.int64),
       array_ops.constant(val2, dtypes.int64),
       array_ops.placeholder_with_default(shape2, shape=[3]))
   sp_tensor3 = sparse_tensor.SparseTensor(
       array_ops.constant([[1, 9], [2, 2], [2, 10]], dtypes.int64),
       array_ops.constant([7, 15, 2], dtypes.int64),
       array_ops.constant([5, 12], dtypes.int64)
   )
   self.sp_tensor3_expected = sparse_tensor.SparseTensorValue(
       [[0, 1, 9], [0, 2, 2], [0, 2, 10], [1, 1, 9], [1, 2, 2], [1, 2, 10]],
       [7, 15, 2, 7, 15, 2],
       [2, 5, 12]
   )
   self.batch_size = 2
   self.key = string_ops.string_join([
       "key_", string_ops.as_string(
           math_ops.cast(10000 * random_ops.random_uniform(()), dtypes.int32))
   ])
   self.sequences = {
       "seq1": np.random.rand(self.value_length, 5),
       "seq2": np.random.rand(self.value_length, 4, 2),
       "seq3": sp_tensor1,
       "seq4": sp_tensor2}
   self.context = {
       "context1": [3, 4],
       "sp_context": sp_tensor3}
   self.initial_states = {
       "state1": np.random.rand(6, 7),
       "state2": np.random.rand(8)
   }

 def testDynamicShapes(self):
   for dtype in [dtypes.float32, dtypes.float64]:
     default_x1 = constant_op.constant(0.1, dtype=dtype)
     default_x2 = constant_op.constant(3.1, dtype=dtype)
     x1 = array_ops.placeholder_with_default(default_x1, shape=None)
     x2 = array_ops.placeholder_with_default(default_x2, shape=None)
     dx1, dx2 = self._nextafter_gradient(x1, x2)
     expected_dx1 = constant_op.constant(1, dtype=dtype)
     expected_dx2 = constant_op.constant(0, dtype=dtype)
     self.assertAllClose(expected_dx1, dx1)
     self.assertAllClose(expected_dx2, dx2)

 def testSampleProbConsistentBroadcastBoth(self):
   with self.test_session() as sess:
     pln = poisson_lognormal.PoissonLogNormalQuadratureCompound(
         loc=array_ops.placeholder_with_default(
             [[0.], [-0.5]],
             shape=[2, 1] if self.static_shape else None),
         scale=array_ops.placeholder_with_default(
             [[1., 0.9]],
             shape=[1, 2] if self.static_shape else None),
         quadrature_size=10,
         validate_args=True)
     self.run_test_sample_consistent_log_prob(
         sess.run, pln, batch_size=4, rtol=0.1, atol=0.08)

 def testMeanVarianceBroadcastScalar(self):
   with self.test_session() as sess:
     pln = poisson_lognormal.PoissonLogNormalQuadratureCompound(
         loc=array_ops.placeholder_with_default(
             [0., -0.5],
             shape=[2] if self.static_shape else None),
         scale=array_ops.placeholder_with_default(
             1.,
             shape=[] if self.static_shape else None),
         quadrature_size=10,
         validate_args=True)
     self.run_test_sample_consistent_mean_variance(
         sess.run, pln, rtol=0.1, atol=0.01)

 def testRaggedTensorSplitsMismatchErrorAtRuntime(self):
   splits1 = array_ops.placeholder_with_default(
       constant_op.constant([0, 3, 3, 5], dtypes.int64), None)
   splits2 = array_ops.placeholder_with_default(
       constant_op.constant([0, 1, 3, 5], dtypes.int64), None)
   x = ragged.from_row_splits([3, 1, 4, 1, 5], splits1)
   y = ragged.from_row_splits([1, 2, 3, 4, 5], splits2)
   result = ragged.map_inner_values(math_ops.add, x, y)
   with self.test_session():
     self.assertRaisesRegexp(
         errors.InvalidArgumentError,
         r'\[Inputs must have identical ragged splits\] '
         r'\[Condition x == y did not hold element-wise:\].*', result.eval)

 def testSampleProbConsistent(self):
   with self.test_session() as sess:
     pln = poisson_lognormal.PoissonLogNormalQuadratureCompound(
         loc=array_ops.placeholder_with_default(
             -2.,
             shape=[] if self.static_shape else None),
         scale=array_ops.placeholder_with_default(
             1.1,
             shape=[] if self.static_shape else None),
         quadrature_size=10,
         validate_args=True)
     self.run_test_sample_consistent_log_prob(
         sess.run, pln, batch_size=1, rtol=0.1)

  def make_mvn(self, dims, new_batch_shape, old_batch_shape):
    new_batch_shape_ph = (
        constant_op.constant(np.int32(new_batch_shape)) if self.is_static_shape
        else array_ops.placeholder_with_default(
            np.int32(new_batch_shape), shape=None))

    scale = np.ones(old_batch_shape + [dims], self.dtype)
    scale_ph = array_ops.placeholder_with_default(
        scale, shape=scale.shape if self.is_static_shape else None)
    mvn = mvn_lib.MultivariateNormalDiag(scale_diag=scale_ph)
    reshape_mvn = batch_reshape_lib.BatchReshape(
        distribution=mvn,
        batch_shape=new_batch_shape_ph,
        validate_args=True)
    return mvn, reshape_mvn

  def operator_and_matrix(
      self, build_info, dtype, use_placeholder,
      ensure_self_adjoint_and_pd=False):
    shape = list(build_info.shape)
    row = np.random.uniform(low=1., high=5., size=shape[:-1])
    col = np.random.uniform(low=1., high=5., size=shape[:-1])

    # Make sure first entry is the same
    row[..., 0] = col[..., 0]

    if ensure_self_adjoint_and_pd:
      # Note that a Toeplitz matrix generated from a linearly decreasing
      # non-negative sequence is positive definite. See
      # https://www.math.cinvestav.mx/~grudsky/Papers/118_29062012_Albrecht.pdf
      # for details.
      row = np.linspace(start=10., stop=1., num=shape[-1])

      # The entries for the first row and column should be the same to guarantee
      # symmetric.
      row = col

    lin_op_row = math_ops.cast(row, dtype=dtype)
    lin_op_col = math_ops.cast(col, dtype=dtype)

    if use_placeholder:
      lin_op_row = array_ops.placeholder_with_default(
          lin_op_row, shape=None)
      lin_op_col = array_ops.placeholder_with_default(
          lin_op_col, shape=None)

    operator = linear_operator_toeplitz.LinearOperatorToeplitz(
        row=lin_op_row,
        col=lin_op_col,
        is_self_adjoint=True if ensure_self_adjoint_and_pd else None,
        is_positive_definite=True if ensure_self_adjoint_and_pd else None)

    flattened_row = np.reshape(row, (-1, shape[-1]))
    flattened_col = np.reshape(col, (-1, shape[-1]))
    flattened_toeplitz = np.zeros(
        [flattened_row.shape[0], shape[-1], shape[-1]])
    for i in range(flattened_row.shape[0]):
      flattened_toeplitz[i] = scipy.linalg.toeplitz(
          flattened_col[i],
          flattened_row[i])
    matrix = np.reshape(flattened_toeplitz, shape)
    matrix = math_ops.cast(matrix, dtype=dtype)

    return operator, matrix

 def test_placeholder_with_default_fed(self):
   with self.test_session() as sess, self.test_scope():
     v = resource_variable_ops.ResourceVariable(4.0)
     ph = array_ops.placeholder_with_default(v, shape=[])
     out = ph * 2
     sess.run(variables.variables_initializer([v]))
     self.assertEqual(2.0, sess.run(out, {ph: 1.0}))

 def _serving_input_receiver_fn():
   """A receiver function to be passed to export_savedmodel."""
   placeholders = {}
   placeholders[feature_keys.TrainEvalFeatures.TIMES] = (
       array_ops.placeholder(
           name=feature_keys.TrainEvalFeatures.TIMES,
           dtype=dtypes.int64,
           shape=[default_batch_size, default_series_length]))
   # Values are only necessary when filtering. For prediction the default
   # value will be ignored.
   placeholders[feature_keys.TrainEvalFeatures.VALUES] = (
       array_ops.placeholder_with_default(
           name=feature_keys.TrainEvalFeatures.VALUES,
           input=array_ops.zeros(
               shape=[
                   default_batch_size
                   if default_batch_size else 0, default_series_length
                   if default_series_length else 0, self._model.num_features
               ],
               dtype=self._model.dtype),
           shape=(default_batch_size, default_series_length,
                  self._model.num_features)))
   if self._model.exogenous_feature_columns:
     with ops.Graph().as_default():
       # Default placeholders have only an unknown batch dimension. Make them
       # in a separate graph, then splice in the series length to the shapes
       # and re-create them in the outer graph.
       parsed_features = (
           feature_column.make_parse_example_spec(
               self._model.exogenous_feature_columns))
       placeholder_features = parsing_ops.parse_example(
           serialized=array_ops.placeholder(
               shape=[None], dtype=dtypes.string),
           features=parsed_features)
       exogenous_feature_shapes = {
           key: (value.get_shape(), value.dtype) for key, value
           in placeholder_features.items()}
     for feature_key, (batch_only_feature_shape, value_dtype) in (
         exogenous_feature_shapes.items()):
       batch_only_feature_shape = (
           batch_only_feature_shape.with_rank_at_least(1).as_list())
       feature_shape = ([default_batch_size, default_series_length]
                        + batch_only_feature_shape[1:])
       placeholders[feature_key] = array_ops.placeholder(
           dtype=value_dtype, name=feature_key, shape=feature_shape)
   # Models may not know the shape of their state without creating some
   # variables/ops. Avoid polluting the default graph by making a new one. We
   # use only static metadata from the returned Tensors.
   with ops.Graph().as_default():
     self._model.initialize_graph()
     model_start_state = self._model.get_start_state()
   for prefixed_state_name, state_tensor in ts_head_lib.state_to_dictionary(
       model_start_state).items():
     state_shape_with_batch = tensor_shape.TensorShape(
         (default_batch_size,)).concatenate(state_tensor.get_shape())
     placeholders[prefixed_state_name] = array_ops.placeholder(
         name=prefixed_state_name,
         shape=state_shape_with_batch,
         dtype=state_tensor.dtype)
   return export_lib.ServingInputReceiver(placeholders, placeholders)

  def make_normal(self, new_batch_shape, old_batch_shape):
    new_batch_shape_ph = (
        constant_op.constant(np.int32(new_batch_shape)) if self.is_static_shape
        else array_ops.placeholder_with_default(
            np.int32(new_batch_shape), shape=None))

    scale = self.dtype(0.5 + np.arange(
        np.prod(old_batch_shape)).reshape(old_batch_shape))
    scale_ph = array_ops.placeholder_with_default(
        scale, shape=scale.shape if self.is_static_shape else None)
    normal = normal_lib.Normal(loc=self.dtype(0), scale=scale_ph)
    reshape_normal = batch_reshape_lib.BatchReshape(
        distribution=normal,
        batch_shape=new_batch_shape_ph,
        validate_args=True)
    return normal, reshape_normal