def _testPad(self, inputs, block_shape, paddings, outputs):
   block_shape = np.array(block_shape)
   paddings = np.array(paddings).reshape((len(block_shape), 2))
   with self.test_session() as sess, self.test_scope():
     for dtype in self.float_types:
       # TODO(b/68813416): Skip bfloat16's as the input type for direct is
       # float32 and results in a mismatch, while making testDirect provide the
       # correctly typed input results in 'no fill-function for data-type'
       # error.
       if dtype == dtypes.bfloat16.as_numpy_dtype:
         continue
       if dtype == np.float16:
         actual_inputs = np.array(inputs).astype(dtype)
         actual_paddings = np.array(paddings).astype(dtype)
         expected_outputs = np.array(outputs).astype(dtype)
       else:
         actual_inputs = inputs
         actual_paddings = paddings
         expected_outputs = outputs
       placeholder = array_ops.placeholder(dtype)
       # outputs = space_to_batch(inputs)
       x_tf = array_ops.space_to_batch_nd(placeholder, block_shape,
                                          actual_paddings)
       self.assertAllEqual(
           sess.run(x_tf, {placeholder: actual_inputs}), expected_outputs)
       # inputs = batch_to_space(outputs)
       placeholder = array_ops.placeholder(dtype)
       x_tf = array_ops.batch_to_space_nd(placeholder, block_shape,
                                          actual_paddings)
       self.assertAllEqual(
           sess.run(x_tf, {placeholder: expected_outputs}), actual_inputs)

  def _testStaticShape(self, input_shape, block_shape, paddings, error):
    block_shape = np.array(block_shape)
    paddings = np.array(paddings)

    # Try with sizes known at graph construction time.
    with self.assertRaises(error):
      _ = array_ops.space_to_batch_nd(
          np.zeros(input_shape, np.float32), block_shape, paddings)

 def _testPad(self, inputs, block_shape, paddings, outputs):
   block_shape = np.array(block_shape)
   paddings = np.array(paddings).reshape((len(block_shape), 2))
   for use_gpu in [False, True]:
     with self.test_session(use_gpu=use_gpu):
       # outputs = space_to_batch(inputs)
       x_tf = array_ops.space_to_batch_nd(
           math_ops.to_float(inputs), block_shape, paddings)
       self.assertAllEqual(x_tf.eval(), outputs)
       # inputs = batch_to_space(outputs)
       x_tf = array_ops.batch_to_space_nd(
           math_ops.to_float(outputs), block_shape, paddings)
       self.assertAllEqual(x_tf.eval(), inputs)

 def _testPad(self, inputs, block_shape, paddings, outputs):
   block_shape = np.array(block_shape)
   paddings = np.array(paddings).reshape((len(block_shape), 2))
   with self.test_session() as sess, self.test_scope():
     for dtype in self.float_types:
       placeholder = array_ops.placeholder(dtype)
       # outputs = space_to_batch(inputs)
       x_tf = array_ops.space_to_batch_nd(placeholder, block_shape, paddings)
       self.assertAllEqual(sess.run(x_tf, {placeholder: inputs}), outputs)
       # inputs = batch_to_space(outputs)
       placeholder = array_ops.placeholder(dtype)
       x_tf = array_ops.batch_to_space_nd(placeholder, block_shape, paddings)
       self.assertAllEqual(sess.run(x_tf, {placeholder: outputs}), inputs)

  def _checkGrad(self, x, block_shape, paddings):
    block_shape = np.array(block_shape)
    paddings = np.array(paddings).reshape((len(block_shape), 2))
    with self.test_session():
      tf_x = ops.convert_to_tensor(x)
      tf_y = array_ops.space_to_batch_nd(tf_x, block_shape, paddings)
      epsilon = 1e-5
      ((x_jacob_t, x_jacob_n)) = gradient_checker.compute_gradient(
          tf_x,
          x.shape,
          tf_y,
          tf_y.get_shape().as_list(),
          x_init_value=x,
          delta=epsilon)

    self.assertAllClose(x_jacob_t, x_jacob_n, rtol=1e-2, atol=epsilon)

  def _testDynamicShape(self, input_shape, block_shape, paddings):
    block_shape = np.array(block_shape)
    paddings = np.array(paddings)
    # Try with sizes unknown at graph construction time.
    input_placeholder = array_ops.placeholder(dtypes.float32)
    block_shape_placeholder = array_ops.placeholder(
        dtypes.int32, shape=block_shape.shape)
    paddings_placeholder = array_ops.placeholder(dtypes.int32)
    t = array_ops.space_to_batch_nd(input_placeholder, block_shape_placeholder,
                                    paddings_placeholder)

    with self.assertRaises(ValueError):
      _ = t.eval({
          input_placeholder: np.zeros(input_shape, np.float32),
          block_shape_placeholder: block_shape,
          paddings_placeholder: paddings
      })

  def testUnknown(self):
    # Verify that input shape and paddings shape can be unknown.
    _ = array_ops.space_to_batch_nd(
        array_ops.placeholder(dtypes.float32),
        array_ops.placeholder(
            dtypes.int32, shape=(2,)),
        array_ops.placeholder(dtypes.int32))

    # Only number of input dimensions is known.
    t = array_ops.space_to_batch_nd(
        array_ops.placeholder(
            dtypes.float32, shape=(None, None, None, None)),
        array_ops.placeholder(
            dtypes.int32, shape=(2,)),
        array_ops.placeholder(dtypes.int32))
    self.assertEqual(4, t.get_shape().ndims)

    # Dimensions are partially known.
    t = array_ops.space_to_batch_nd(
        array_ops.placeholder(
            dtypes.float32, shape=(None, None, None, 2)),
        array_ops.placeholder(
            dtypes.int32, shape=(2,)),
        array_ops.placeholder(dtypes.int32))
    self.assertEqual([None, None, None, 2], t.get_shape().as_list())

    # Dimensions are partially known.
    t = array_ops.space_to_batch_nd(
        array_ops.placeholder(
            dtypes.float32, shape=(3, None, None, 2)), [2, 3],
        array_ops.placeholder(dtypes.int32))
    self.assertEqual([3 * 2 * 3, None, None, 2], t.get_shape().as_list())

    # Dimensions are partially known.
    t = array_ops.space_to_batch_nd(
        array_ops.placeholder(
            dtypes.float32, shape=(3, None, 2, 2)), [2, 3], [[1, 1], [0, 1]])
    self.assertEqual([3 * 2 * 3, None, 1, 2], t.get_shape().as_list())

    # Dimensions are fully known.
    t = array_ops.space_to_batch_nd(
        array_ops.placeholder(
            dtypes.float32, shape=(3, 2, 3, 2)), [2, 3], [[1, 1], [0, 0]])
    self.assertEqual([3 * 2 * 3, 2, 1, 2], t.get_shape().as_list())

def _BatchToSpaceNDGrad(op, grad):
  # Its gradient is the opposite op: SpaceToBatchND.
  return [array_ops.space_to_batch_nd(grad, op.inputs[1], op.inputs[2]),
          None, None]