def testCaching(self):
    """Confirm caching of control output is recalculated between calls."""
    a = constant_op.constant(1)
    b = constant_op.constant(2)
    with ops.control_dependencies([a]):
      c = constant_op.constant(42)

    shared = {}

    def sub(t):
      shared[t] = shared.get(t, 0) + 1
      return t

    a = subscribe.subscribe(a,
                            lambda t: script_ops.py_func(sub, [t], [t.dtype]))

    with ops.control_dependencies([b]):
      d = constant_op.constant(11)

    # If it was using outdated cached control_outputs then
    # evaling would not trigger the new subscription.
    b = subscribe.subscribe(b,
                            lambda t: script_ops.py_func(sub, [t], [t.dtype]))

    with self.cached_session() as sess:
      c_out = self.evaluate([c])
      d_out = self.evaluate([d])

    self.assertEqual(c_out, [42])
    self.assertEqual(d_out, [11])
    self.assertEqual(shared, {2: 1, 1: 1})

  def test_subscribe_tensors_on_different_devices(self):
    """Side effect ops are added with the same device of the subscribed op."""
    c1 = constant_op.constant(10)
    c2 = constant_op.constant(20)

    with ops.device('cpu:0'):
      add = math_ops.add(c1, c2)

    with ops.device('cpu:1'):
      mul = math_ops.multiply(c1, c2)

    def sub(t):
      return t

    add_sub = subscribe.subscribe(
        add, lambda t: script_ops.py_func(sub, [t], [t.dtype]))

    mul_sub = subscribe.subscribe(
        mul, lambda t: script_ops.py_func(sub, [t], [t.dtype]))

    # Expect the identity tensors injected by subscribe to have been created
    # on the same device as their original tensors.
    self.assertNotEqual(add_sub.device, mul_sub.device)
    self.assertEqual(add.device, add_sub.device)
    self.assertEqual(mul.device, mul_sub.device)

 def testLarge(self):
   with self.test_session() as sess:
     x = array_ops.zeros([1000000], dtype=np.float32)
     y = script_ops.py_func(lambda x: x + 1, [x], [dtypes.float32])
     z = script_ops.py_func(lambda x: x * 2, [x], [dtypes.float32])
     for _ in xrange(100):
       sess.run([y[0].op, z[0].op])

 def testGradientFunction(self):
   # Input to tf.py_func is necessary, otherwise get_gradient_function()
   # returns None per default.
   a = constant_op.constant(0)
   x, = script_ops.py_func(lambda a: 0, [a], [dtypes.int64])
   y, = script_ops.py_func(lambda a: 0, [a], [dtypes.int64], stateful=False)
   self.assertEqual(None, ops.get_gradient_function(x.op))
   self.assertEqual(None, ops.get_gradient_function(y.op))

 def make_graphs():
   for _ in xrange(1000):
     g = ops.Graph()
     with g.as_default():
       c = constant_op.constant([1.], dtypes.float32)
       _ = script_ops.py_func(lambda x: x + 1, [c], [dtypes.float32])
       _ = script_ops.eager_py_func(lambda x: x + 1, [c], [dtypes.float32])
       # These ops have a reference to 'c' which has a reference to the graph.
       # Checks if the functions are being deleted though the graph is referenced from them.
       # (see #18292)
       _ = script_ops.py_func(lambda x: x + c.shape[0], [c], [dtypes.float32])
       _ = script_ops.eager_py_func(lambda x: x + c.shape[0], [c], [dtypes.float32])

    def testStrings(self):
        def read_fixed_length_numpy_strings():
            return np.array([b" there"])

        def read_and_return_strings(x, y):
            return x + y

        with self.test_session():
            x = constant_op.constant([b"hello", b"hi"], dtypes.string)
            y, = script_ops.py_func(read_fixed_length_numpy_strings, [], [dtypes.string])
            z, = script_ops.py_func(read_and_return_strings, [x, y], [dtypes.string])
            self.assertListEqual(list(z.eval()), [b"hello there", b"hi there"])

  def testObjectArraysAreConvertedToBytes(self):

    def read_object_array():
      return np.array([b" there", u" ya"], dtype=np.object)

    def read_and_return_strings(x, y):
      return x + y

    with self.test_session():
      x = constant_op.constant(["hello", "hi"], dtypes.string)
      y, = script_ops.py_func(read_object_array, [],
                              [dtypes.string])
      z, = script_ops.py_func(read_and_return_strings, [x, y], [dtypes.string])
      self.assertListEqual(list(z.eval()), [b"hello there", b"hi ya"])

  def testSideEffect(self):
    a = constant_op.constant(1)
    b = constant_op.constant(1)
    c = math_ops.add(a, b)
    with ops.control_dependencies([c]):
      d = constant_op.constant(42)
    n = math_ops.negative(c)

    shared = []

    def sub(t):
      shared.append(t)
      return t

    c0 = c
    self.assertTrue(c0.op in d.op.control_inputs)
    c = subscribe.subscribe(c,
                            lambda t: script_ops.py_func(sub, [t], [t.dtype]))
    # Verify that control dependencies are correctly moved to the subscription.
    self.assertFalse(c0.op in d.op.control_inputs)
    self.assertTrue(c.op in d.op.control_inputs)

    with self.cached_session() as sess:
      c_out = self.evaluate([c])
      n_out = self.evaluate([n])
      d_out = self.evaluate([d])

    self.assertEqual(n_out, [-2])
    self.assertEqual(c_out, [2])
    self.assertEqual(d_out, [42])
    self.assertEqual(shared, [2, 2, 2])

  def testSideEffect(self):
    a = constant_op.constant(1)
    b = constant_op.constant(1)
    c = math_ops.add(a, b)
    with ops.control_dependencies([c]):
      d = constant_op.constant(42)
    n = math_ops.negative(c)

    shared = []

    def sub(t):
      shared.append(t)
      return t

    c = subscribe.subscribe(c,
                            lambda t: script_ops.py_func(sub, [t], [t.dtype]))

    with self.test_session() as sess:
      c_out = sess.run([c])
      n_out = sess.run([n])
      d_out = sess.run([d])

    self.assertEquals(n_out, [-2])
    self.assertEquals(c_out, [2])
    self.assertEquals(d_out, [42])
    self.assertEquals(shared, [2, 2, 2])

  def testMultipleOutputs(self):
    """Handle subscriptions to multiple outputs from the same op."""
    sparse_tensor_1 = sparse_tensor.SparseTensor(
        indices=[[0, 0], [1, 2]], values=[1, 2], dense_shape=[3, 4])
    sparse_tensor_2 = sparse_tensor.SparseTensor(
        indices=[[0, 0], [1, 2]], values=[2, 3], dense_shape=[3, 4])

    # This op has three outputs.
    sparse_add = sparse_ops.sparse_add(sparse_tensor_1, sparse_tensor_2)

    self.assertEqual(3, len(sparse_add.op.outputs))

    c1 = constant_op.constant(1)

    with ops.control_dependencies(sparse_add.op.outputs):
      # This op depends on all the three outputs.
      neg = -c1

    shared = []
    def sub(t):
      shared.append(t)
      return t

    # Subscribe the three outputs at once.
    subscribe.subscribe(sparse_add.op.outputs,
                        lambda t: script_ops.py_func(sub, [t], [t.dtype]))

    with self.cached_session() as sess:
      self.evaluate([neg])

    # All three ops have been processed.
    self.assertEqual(3, len(shared))

  def testResourceType(self):
    """Confirm that subscribe correctly handles tensors with 'resource' type."""
    tensor_array = tensor_array_ops.TensorArray(
        dtype=dtypes.float32,
        tensor_array_name='test',
        size=3,
        infer_shape=False)
    writer = tensor_array.write(0, [[4.0, 5.0]])
    reader = writer.read(0)

    shared = []

    def sub(t):
      shared.append(t)
      return t

    # TensorArray's handle output tensor has a 'resource' type and cannot be
    # subscribed as it's not 'numpy compatible' (see dtypes.py).
    # Expect that the original tensor is returned when subscribing to it.
    tensor_array_sub = subscribe.subscribe(
        tensor_array.handle, lambda t: script_ops.py_func(sub, [t], [t.dtype]))
    self.assertIs(tensor_array_sub, tensor_array.handle)
    self.assertFalse(subscribe._is_subscribed_identity(tensor_array.handle))

    with self.cached_session() as sess:
      self.evaluate([reader])
    self.assertEqual(0, len(shared))

  def testMapAndBatchOutOfRangeError(self, threshold):

    def raising_py_fn(i):
      if i >= threshold:
        raise StopIteration()
      else:
        return i

    iterator = (
        dataset_ops.Dataset.range(100).apply(
            batching.map_and_batch(
                lambda x: script_ops.py_func(raising_py_fn, [x], dtypes.int64),
                batch_size=10)).make_one_shot_iterator())
    get_next = iterator.get_next()

    with self.cached_session() as sess:
      for i in range(threshold // 10):
        self.assertAllEqual([i * 10 + j for j in range(10)],
                            self.evaluate(get_next))
      if threshold % 10 != 0:
        self.assertAllEqual(
            [threshold // 10 * 10 + j for j in range(threshold % 10)],
            sess.run(get_next))
      with self.assertRaises(errors.OutOfRangeError):
        sess.run(get_next)

  def testMapAndBatchMapError(self, threshold, numa_aware):

    def raising_py_fn(i):
      if i >= threshold:
        raise StopIteration()
      else:
        return i

    dataset = dataset_ops.Dataset.range(100).apply(
        batching.map_and_batch(
            lambda x: script_ops.py_func(raising_py_fn, [x], dtypes.int64),
            batch_size=10))
    if numa_aware:
      options = dataset_ops.Options()
      options.experimental_numa_aware = True
      dataset = dataset.with_options(options)

    get_next = self.getNext(dataset)
    for i in range(threshold // 10):
      self.assertAllEqual([i * 10 + j for j in range(10)],
                          self.evaluate(get_next()))
    if numa_aware:
      if threshold % 10 != 0:
        self.assertAllEqual(
            [threshold // 10 * 10 + j for j in range(threshold % 10)],
            self.evaluate(get_next()))
    else:
      for i in range(threshold // 10, 10):
        with self.assertRaises(errors.InvalidArgumentError):
          self.evaluate(get_next())
    with self.assertRaises(errors.OutOfRangeError):
      self.evaluate(get_next())

 def testAlias(self):
   with self.test_session():
     np_array = np.array([1.0, 2.0], dtype=np.float32)
     tf_array = script_ops.py_func(lambda: np_array, [], [dtypes.float32])
     value = tf_array + constant_op.constant([2.0, 3.0], dtype=dtypes.float32)
     value.op.run()
     self.assertAllEqual(np_array, [1.0, 2.0])

  def _testNumThreadsHelper(self, num_threads, override_threadpool_fn):

    def get_thread_id(_):
      # Python creates a dummy thread object to represent the current
      # thread when called from an "alien" thread (such as a
      # `PrivateThreadPool` thread in this case). It does not include
      # the TensorFlow-given display name, but it has a unique
      # identifier that maps one-to-one with the underlying OS thread.
      return np.array(threading.current_thread().ident).astype(np.int64)

    dataset = (
        dataset_ops.Dataset.range(1000).map(
            lambda x: script_ops.py_func(get_thread_id, [x], dtypes.int64),
            num_parallel_calls=32).apply(unique.unique()))
    dataset = override_threadpool_fn(dataset)
    next_element = self.getNext(dataset, requires_initialization=True)

    thread_ids = []
    try:
      while True:
        thread_ids.append(self.evaluate(next_element()))
    except errors.OutOfRangeError:
      pass
    self.assertLen(thread_ids, len(set(thread_ids)))
    self.assertNotEmpty(thread_ids)
    if num_threads:
      # NOTE(mrry): We don't control the thread pool scheduling, and
      # so cannot guarantee that all of the threads in the pool will
      # perform work.
      self.assertLessEqual(len(thread_ids), num_threads)

  def testMapAndBatchOutOfRangeError(self, threshold, numa_aware):

    def raising_py_fn(i):
      if i == threshold:
        raise StopIteration()
      elif i > threshold:
        raise RuntimeError("Alternate error; you shouldn't see me! (i: %s)" % i)
      else:
        return i

    dataset = dataset_ops.Dataset.range(100).apply(
        batching.map_and_batch(
            lambda x: script_ops.py_func(raising_py_fn, [x], dtypes.int64),
            batch_size=10))
    if numa_aware:
      options = dataset_ops.Options()
      options.experimental_numa_aware = True
      dataset = dataset.with_options(options)
    iterator = dataset_ops.make_one_shot_iterator(dataset)
    get_next = iterator.get_next()

    with self.cached_session() as sess:
      for i in range(threshold // 10):
        self.assertAllEqual([i * 10 + j for j in range(10)],
                            self.evaluate(get_next))
      if threshold % 10 != 0:
        self.assertAllEqual(
            [threshold // 10 * 10 + j for j in range(threshold % 10)],
            self.evaluate(get_next))
      with self.assertRaises(errors.OutOfRangeError):
        self.evaluate(get_next)

  def testOverrideThreadPool(self):

    def get_thread_id(_):
      # Python creates a dummy thread object to represent the current
      # thread when called from an "alien" thread (such as a
      # `PrivateThreadPool` thread in this case). It does not include
      # the TensorFlow-given display name, but it has a unique
      # identifier that maps one-to-one with the underlying OS thread.
      return np.array(threading.current_thread().ident).astype(np.int64)

    for num_threads in [1, 2, 4, 8, 16]:

      dataset = (
          Dataset.range(1000).map(
              lambda x: script_ops.py_func(get_thread_id, [x], dtypes.int64),
              num_parallel_calls=32).apply(unique.unique()))

      dataset = threadpool.override_threadpool(
          dataset,
          threadpool.PrivateThreadPool(
              num_threads, display_name='private_thread_pool_%d' % num_threads))

      thread_ids = []
      for next_element in datasets.Iterator(dataset):
        thread_ids.append(next_element)
      self.assertEqual(len(thread_ids), len(set(thread_ids)))
      self.assertGreater(len(thread_ids), 0)
      # NOTE(mrry): We don't control the thread pool scheduling, and
      # so cannot guarantee that all of the threads in the pool will
      # perform work.
      self.assertLessEqual(len(thread_ids), num_threads)

def _batch_to_patches(batch, patches_per_image, patch_size):
  """Extract patches from a batch.

  Args:
      batch: (tensor) The batch of images (batch, height, width, channels).
      patches_per_image: (int) Number of patches to extract per image.
      patch_size: (int) Size of the patches (size, size, channels) to extract.
  Returns:
      Tensor (batch*patches_per_image, patch_size, patch_size, channels) of
      patches.
  """

  def py_func_random_patches(batch):
    """Numpy wrapper."""
    batch_size, height, width, channels = batch.shape
    patch_count = patches_per_image * batch_size
    hs = patch_size // 2
    # Randomly pick patches.
    patch_id, y, x, chan = np.ogrid[0:patch_count, -hs:hs + 1, -hs:hs + 1, 0:3]
    img_id = patch_id // patches_per_image
    # pylint: disable=g-no-augmented-assignment
    # Need explicit addition for broadcast to work properly.
    y = y + np.random.randint(hs, height - hs, size=(patch_count, 1, 1, 1))
    x = x + np.random.randint(hs, width - hs, size=(patch_count, 1, 1, 1))
    # pylint: enable=g-no-augmented-assignment
    idx = ((img_id * height + y) * width + x) * channels + chan
    patches = batch.flat[idx]
    return patches

  patches = script_ops.py_func(
      py_func_random_patches, [batch], batch.dtype, stateful=False)
  return patches

    def generator_map_fn(iterator_id_t):
      """Generates the next element from iterator with ID `iterator_id_t`.

      We map this function across an infinite repetition of the
      `iterator_id_t`, and raise `StopIteration` to terminate the iteration.

      Args:
        iterator_id_t: A `tf.int64` tensor whose value uniquely identifies
          the iterator in `generator_state` from which to generate an element.

      Returns:
        A nested structure of tensors representing an element from the iterator.
      """

      def generator_py_func(iterator_id):
        """A `py_func` that will be called to invoke the iterator."""
        try:
          values = next(generator_state.get_iterator(iterator_id))
        except StopIteration:
          generator_state.iterator_completed(iterator_id)
          raise StopIteration("Iteration finished.")

        # Use the same _convert function from the py_func() implementation to
        # convert the returned values to arrays early, so that we can inspect
        # their values.
        # pylint: disable=protected-access
        ret_arrays = [
            script_ops.FuncRegistry._convert(ret, dtype=dtype.as_numpy_dtype)
            for ret, dtype in zip(nest.flatten_up_to(output_types, values),
                                  flattened_types)
        ]
        # pylint: enable=protected-access

        # Additional type and shape checking to ensure that the components
        # of the generated element match the `output_types` and `output_shapes`
        # arguments.
        for (ret_array, expected_dtype, expected_shape) in zip(
            ret_arrays, flattened_types, flattened_shapes):
          if ret_array.dtype != expected_dtype.as_numpy_dtype:
            raise TypeError(
                "`generator` yielded an element of type %s where an element "
                "of type %s was expected." % (ret_array.dtype,
                                              expected_dtype.as_numpy_dtype))
          if not expected_shape.is_compatible_with(ret_array.shape):
            raise ValueError(
                "`generator` yielded an element of shape %s where an element "
                "of shape %s was expected." % (ret_array.shape, expected_shape))

        return ret_arrays

      flat_values = script_ops.py_func(
          generator_py_func, [iterator_id_t], flattened_types, stateful=True)

      # The `py_func()` op drops the inferred shapes, so we add them back in
      # here.
      if output_shapes is not None:
        for ret_t, shape in zip(flat_values, flattened_shapes):
          ret_t.set_shape(shape)

      return nest.pack_sequence_as(output_types, flat_values)

 def testArray(self):
   with self.test_session():
     x = constant_op.constant([1.0, 2.0], dtypes.float64)
     y = constant_op.constant([2.0, 3.0], dtypes.float64)
     z = self.evaluate(script_ops.py_func(np_func, [x, y], [dtypes.float64]))
     self.assertAllEqual(z[0],
                         np_func([1.0, 2.0], [2.0, 3.0]).astype(np.float64))