def testFeedTwoHandlesDirectly(self):
    with self.test_session() as sess:
      a = constant_op.constant(10.0)
      b = constant_op.constant(5.0)
      c = math_ops.multiply(a, b)
      d = math_ops.div(a, b)
      e = math_ops.subtract(c, d)

      h_c = sess.run(session_ops.get_session_handle(c))
      h_d = sess.run(session_ops.get_session_handle(d))

      self.assertAllClose(48.0, sess.run(e, feed_dict={c: h_c, d: h_d}))
      self.assertAllClose(-48.0, sess.run(e, feed_dict={c: h_d, d: h_c}))

  def testHandleForLoop(self):
    with self.test_session() as sess:
      # Initialize a handle.
      a = constant_op.constant(0)
      h = session_ops.get_session_handle(a)
      h = sess.run(h)

      # Do some computation.
      f, x = session_ops.get_session_tensor(h.handle, dtypes.int32)
      # Must define the loop body outside the loop.
      h_x = session_ops.get_session_handle(math_ops.add(x, 1))
      for _ in range(100):
        # This exercises garbage collection.
        h = sess.run(h_x, feed_dict={f: h.handle})

      self.assertEqual(100, h.eval())

  def testDirectHandleFeedOverlappingWithFetches(self):
    with self.cached_session() as sess:
      a = constant_op.constant(10.0)
      b = constant_op.constant(5.0)
      c = math_ops.multiply(a, b)
      h_c = sess.run(session_ops.get_session_handle(c))
      d = array_ops.identity(c)

      c_val = sess.run(c, feed_dict={c: h_c})
      self.assertAllClose(50.0, c_val)

      d_val = sess.run(d, feed_dict={c: h_c})
      self.assertAllClose(50.0, d_val)

      c_val, d_val = sess.run([c, d], feed_dict={c: h_c, d: 60.0})
      self.assertAllClose(50.0, c_val)
      self.assertAllClose(60.0, d_val)

      c_val, d_val = sess.run([c, d], feed_dict={c: 60.0, d: h_c})
      self.assertAllClose(60.0, c_val)
      self.assertAllClose(50.0, d_val)

      c_val, d_val = sess.run([c, d], feed_dict={c: h_c, d: h_c})
      self.assertAllClose(50.0, c_val)
      self.assertAllClose(50.0, d_val)

  def numpy_to_handle(self, array):
    """Upload numpy array into TensorFlow runtime.

    Args:
      array: numpy array to convert to TensorHandle

    Returns:
      TensorHandle corresponding to given numpy array.
    """

    tf_dtype = dtypes.as_dtype(array.dtype)
    current_device = get_current_device_string(self.g)
    current_device_sanitized = current_device.replace(":", "")
    key = ("numpy_to_handle", tf_dtype.name, current_device)

    if key in self.op_cache:
      holder, handle_op = self.op_cache[key]
    else:
      if self.PRINT_CACHE_MISSES:
        print("Imperative cache miss for %s"%(str(key)))

      op_prefix = "numpy_to_handle.%s.%s" % (tf_dtype.name,
                                             current_device_sanitized)
      with self.g.as_default():
        holder = array_ops.placeholder(dtype=array.dtype,
                                       name=op_prefix+".holder")
        handle_op = session_ops.get_session_handle(holder,
                                                   name=op_prefix+".handle")
      self.op_cache[key] = (holder, handle_op)

    handle = self.run(handle_op, feed_dict={holder: array})
    return handle

 def testHandleDelete(self):
   with self.test_session() as sess:
     # Return a handle.
     a = constant_op.constant(10)
     b = constant_op.constant(5)
     c = math_ops.multiply(a, b)
     h = session_ops.get_session_handle(c)
     sess.run(h).delete()

  def tensor_to_itensor(self, tensor):

    op_prefix = "tensor_to_itensor"
    with self.g.as_default():
      handle_op = session_ops.get_session_handle(tensor,
                                                 name=op_prefix+".handle")
    handle = self.run(handle_op)
    return ITensor(self, handle)

  def testMultiDevices(self):
    with self.test_session() as sess:
      with ops.device(test.gpu_device_name()):
        a = constant_op.constant(1.0)
        a_handle = sess.run(session_ops.get_session_handle(a))
      with ops.device("/cpu:0"):
        b = constant_op.constant(2.0)
        b_handle = sess.run(session_ops.get_session_handle(b))

      a_p, a_t = session_ops.get_session_tensor(a_handle.handle, dtypes.float32)
      b_p, b_t = session_ops.get_session_tensor(b_handle.handle, dtypes.float32)
      c = math_ops.add(a_t, b_t)
      c_handle = sess.run(
          session_ops.get_session_handle(c),
          feed_dict={a_p: a_handle.handle,
                     b_p: b_handle.handle})
      self.assertEqual(3.0, c_handle.eval())

  def testHandleWhileLoop(self):
    with self.test_session() as sess:
      # Initialize a handle.
      a = constant_op.constant(0)
      h = session_ops.get_session_handle(a)
      h = sess.run(h)

      # Do some computation.
      f, x = session_ops.get_session_tensor(h.handle, dtypes.int32)
      b = constant_op.constant(100)
      p = math_ops.less(x, b)
      # Must define the loop body outside the loop.
      h_x = session_ops.get_session_handle(math_ops.add(x, 1))
      while True:
        rp, h = sess.run([p, h_x], feed_dict={f: h.handle})
        if not rp:
          break

      self.assertEqual(101, h.eval())

  def make_function(self, inputs, outputs, name=""):
    """Create callable that accept argument ITensors in the same order as
    inputs argument, and produces tuple of outputs which are ITensors
    corresponding to outputs.

    Example usage:
    x0 = env.tf.ones()       # create ITensor
    x = env.make_input(x0) # create Tensor
    y = env.make_input(x0) # create Tensor
    z1 = tf.add(x, y)         # operate on Tensors
    z2 = tf.sub(x, y)         # operate on Tensors
    f = env.make_function(inputs=[x, y], outputs=[z1, z2])

    print(f(x0, x0*5))       # feed ITensors, get result back as ITensors
    """

    input_holders = []
    for input_ in inputs:
      input_holders.append(self.input_dict[input_])

    output_handle_ops = []
    if is_list_or_tuple(outputs):
      for (i,tensor) in enumerate(outputs):
        op_name = "custom_function_%s.output.%s"%(name, i)
        output_handle_ops.append(session_ops.get_session_handle(tensor,
                                                                op_name))
    # special-case single output
    else:
      op_name = "custom_function_%s.output"%(name)
      output_handle_ops = session_ops.get_session_handle(outputs, op_name)

    def func(*args):
      feed_dict = {}
      for (i, arg) in enumerate(args):
        feed_dict[input_holders[i]] = arg.tf_handle

      tensor_handles = self.sess.run(output_handle_ops, feed_dict=feed_dict)
      if is_list_or_tuple(tensor_handles):
        return [ITensor(self, t) for t in tensor_handles]
      else:
        return ITensor(self, tensor_handles)
      
    return func

  def testFeedOneHandleDirectly(self):
    with self.test_session() as sess:
      a = constant_op.constant(10.0)
      b = constant_op.constant(5.0)
      c = math_ops.multiply(a, b)
      d = math_ops.multiply(c, c)

      h_c = sess.run(session_ops.get_session_handle(c))

      self.assertAllClose(2500.0, sess.run(d, feed_dict={c: h_c}))

  def testHandlePlacement(self):
    with self.test_session() as sess:
      a = constant_op.constant(1.0)
      a_handle_op = session_ops.get_session_handle(a)
      b = constant_op.constant(2.0)
      b_handle_op = session_ops.get_session_handle(b)

      a_handle = sess.run(a_handle_op)
      b_handle = sess.run(b_handle_op)

      a_p, a_t = session_ops.get_session_tensor(a_handle.handle, dtypes.float32)
      b_p, b_t = session_ops.get_session_tensor(b_handle.handle, dtypes.float32)

      c = math_ops.add(a_t, b_t)
      c_handle = sess.run(
          session_ops.get_session_handle(c),
          feed_dict={a_p: a_handle.handle,
                     b_p: b_handle.handle})
      self.assertEqual(3.0, c_handle.eval())

  def testHandleEval(self):
    with self.test_session() as sess:
      # Return a handle.
      a = constant_op.constant(10)
      b = constant_op.constant(5)
      c = math_ops.multiply(a, b)
      h = session_ops.get_session_handle(c)
      h = sess.run(h)

      # Get the tensor from its handle.
      self.assertEqual(50, h.eval())

  def testHandleMover(self):
    with self.test_session() as sess:
      # Return a handle.
      a = constant_op.constant(10)
      b = constant_op.constant(5)
      c = math_ops.multiply(a, b)
      h = session_ops.get_session_handle(c)
      h = sess.run(h)

      # Feed a tensor handle.
      f, x = session_ops.get_session_tensor(h.handle, dtypes.int32)
      y = math_ops.multiply(x, 10)
      self.assertEqual(500, sess.run(y, feed_dict={f: h.handle}))

      # Feed another tensor handle.
      with ops.device(test.gpu_device_name()):
        a = constant_op.constant(10)
        h = session_ops.get_session_handle(a)
        h = sess.run(h)
        self.assertEqual(100, sess.run(y, feed_dict={f: h.handle}))

  def testHandleAndValue(self):
    with self.test_session() as sess:
      # Return a handle and a value.
      a = constant_op.constant(10)
      b = constant_op.constant(5)
      c = math_ops.multiply(a, b)
      h = session_ops.get_session_handle(c)
      v = math_ops.multiply(a, c)
      h, v = sess.run([h, v])

      self.assertEqual(50, h.eval())
      self.assertEqual(500, v)

  def testHandleBasic(self):
    with self.test_session() as sess:
      # Return a handle.
      a = constant_op.constant(10)
      b = constant_op.constant(5)
      c = math_ops.multiply(a, b)
      h = session_ops.get_session_handle(c)
      h = sess.run(h)

      # Feed a tensor handle.
      f, x = session_ops.get_session_tensor(h.handle, dtypes.int32)
      y = math_ops.multiply(x, 10)
      self.assertEqual(500, sess.run(y, feed_dict={f: h.handle}))

  def testFeedHandleToVariableDirectly(self):
    with self.test_session() as sess:
      a = variables.Variable(12.0)
      inc_a = state_ops.assign_add(a, 2.0)
      b = math_ops.add(a, 5.0)
      sess.run(a.initializer)

      h_a_read = sess.run(session_ops.get_session_handle(a.read_value()))
      self.assertAllClose(12.0, sess.run(a))

      self.assertAllClose(17.0, sess.run(b, feed_dict={a: h_a_read}))
      sess.run(inc_a)
      self.assertAllClose(19.0, sess.run(b, feed_dict={a: h_a_read}))

  def testHandleDeleteRaw(self):
    with self.test_session() as sess:
      # Return a handle.
      a = constant_op.constant(10)
      b = constant_op.constant(5)
      c = math_ops.multiply(a, b)
      h = session_ops.get_session_handle(c)
      h = sess.run(h)

      # Delete using a raw tensor handle.
      raw_h = h.get_raw_handle()
      f, x = session_ops.delete_session_tensor(raw_h)
      sess.run(x, feed_dict={f: raw_h})

  def testHandleGC(self):
    with self.test_session() as sess:
      # initial values live on CPU
      with ops.device("/cpu:0"):
        one = constant_op.constant(1, dtype=dtypes.float32)
        one_handle = sess.run(session_ops.get_session_handle(one))
        x_handle = sess.run(session_ops.get_session_handle(one))

      # addition lives on GPU
      with ops.device(test.gpu_device_name()):
        add_h1, add_t1 = session_ops.get_session_tensor(one_handle.handle,
                                                        dtypes.float32)
        add_h2, add_t2 = session_ops.get_session_tensor(x_handle.handle,
                                                        dtypes.float32)
        add_op = math_ops.add(add_t1, add_t2)
        add_output = session_ops.get_session_handle(add_op)

      # add 1 to tensor 20 times
      for _ in range(20):
        x_handle = sess.run(
            add_output,
            feed_dict={add_h1: one_handle.handle,
                       add_h2: x_handle.handle})

  def testHandleCond(self):
    with self.test_session() as sess:
      # Return a handle and a value
      a = constant_op.constant(10)
      b = constant_op.constant(5)
      p = math_ops.less(a, b)
      c = math_ops.multiply(a, b)
      h = session_ops.get_session_handle(c)
      p, h = sess.run([p, h])

      # Run by feeding a tensor handle.
      f, x = session_ops.get_session_tensor(h.handle, dtypes.int32)
      if p:
        y = math_ops.multiply(x, 10)
      else:
        y = math_ops.multiply(x, 100)
      result = sess.run(y, feed_dict={f: h.handle})

      self.assertEqual(5000, result)

  def sum1(self, input_itensor):
    """Create a specialized op that sums over 1 dimensional vector.
    This avoids having to create Rank/Range ops that initialize indices
    in the default tf.reduce_sum."""

    op_type_name = "sum1"
    tf_dtype = input_itensor.dtype
    current_device = get_current_device_string(self.g)
    current_device_sanitized = current_device.replace(":", "")
    key = (op_type_name, tf_dtype.name, current_device_sanitized)

    if key in self.op_cache:
      if self.PRINT_CACHE_HITS:
        print("Imperative cache hit for %s"%(str(key)))
      op = self.op_cache[key]
    else:
      if self.PRINT_CACHE_MISSES:
        print("Imperative cache miss for %s"%(str(key)))
      with self.g.as_default():
        op_prefix = op_type_name + "." + tf_dtype.name
        holder, tensor = session_ops.get_session_tensor(
            input_itensor.tf_handle, input_itensor.dtype, name=op_prefix+".0")
        input_holders = {"input": holder}
        reduction_indices = constant_op.constant([0], dtype=dtypes.int32,
                                                 name=op_prefix+".1")
        output = gen_math_ops._sum(input=tensor,
                                   reduction_indices=reduction_indices,
                                   keep_dims=False, name=op_prefix+".op")
        op_prefix = op_prefix+".out"
        output_handle = session_ops.get_session_handle(output,
                                                       op_prefix+".handle")

      op = Op(self, input_holders, output_handle)
      self.cache_add(key, op)

    return op(input=input_itensor)