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.batch_to_space_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, crops):
    block_shape = np.array(block_shape)
    crops = np.array(crops).reshape((len(block_shape), 2))
    with self.test_session():
      tf_x = ops.convert_to_tensor(x)
      tf_y = array_ops.batch_to_space_nd(tf_x, block_shape, crops)
      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.batch_to_space_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 _CloneWithNewOperands(layer_op, input_tensor, weight_tensor,
                          batch_to_space_op):
  """Clones layer_op with input_tensor and weight_tensor as new inputs."""
  new_layer_name = layer_op.name.split('/')[-1] + '_Fold'
  if layer_op.type == 'Conv2D':
    return nn_ops.conv2d(
        input_tensor,
        weight_tensor,
        strides=layer_op.get_attr('strides'),
        padding=layer_op.get_attr('padding'),
        use_cudnn_on_gpu=layer_op.get_attr('use_cudnn_on_gpu'),
        data_format=layer_op.get_attr('data_format'),
        name=new_layer_name)
  elif layer_op.type == 'MatMul':
    return math_ops.matmul(
        input_tensor,
        weight_tensor,
        transpose_a=layer_op.get_attr('transpose_a'),
        transpose_b=layer_op.get_attr('transpose_b'),
        name=new_layer_name)
  elif layer_op.type == 'DepthwiseConv2dNative':
    # We don't copy dilation rate because we reuse the input SpaceToBatch
    # and create our own BatchToSpace operation below.
    conv = nn.depthwise_conv2d(
        input_tensor,
        weight_tensor,
        strides=layer_op.get_attr('strides'),
        padding=layer_op.get_attr('padding'),
        name=new_layer_name)
    # Copy the batch to space operation if we have a atrous convolution.
    if batch_to_space_op:
      batch_to_space_op = layer_op.outputs[0].consumers()[0]
      # TODO(suharshs): It's hard to make this name match with the unfused name.
      # Restructure this code to not rely on scope at all.
      new_batch_to_space_name = batch_to_space_op.name.split('/')[-1] + '_Fold'
      conv = array_ops.batch_to_space_nd(
          conv,
          batch_to_space_op.inputs[1],
          batch_to_space_op.inputs[2],
          name=new_batch_to_space_name)
    return conv
  else:
    raise ValueError('Cannot handle operation of type: %s' % layer_op.type)

  def testUnknownShape(self):
    # Verify that input shape and paddings shape can be unknown.
    _ = array_ops.batch_to_space_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.batch_to_space_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.batch_to_space_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.batch_to_space_nd(
        array_ops.placeholder(
            dtypes.float32, shape=(3 * 2 * 3, None, None, 2)), [2, 3],
        array_ops.placeholder(dtypes.int32))
    self.assertEqual([3, None, None, 2], t.get_shape().as_list())

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

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

def _SpaceToBatchNDGrad(op, grad):
  # Its gradient is the opposite op: BatchToSpaceND.
  return [array_ops.batch_to_space_nd(grad, op.inputs[1], op.inputs[2]),
          None, None]

def _CreateFoldedOp(graph, context, has_scaling, freeze_batch_norm_delay,
                    is_training):
  """Folds in batch norm layer into preceding convolution or FC layer.

  Creates 3 new nodes, connects their inputs and adds them to the graph:
  mul is cloned into mul_fold, Conv2D or MatMul, or DepthwiseConv2d is cloned
  into respective *_Fold, add is cloned into add_fold.

  Args:
    graph: Graph to modify.
    context: String, batch norm context, i.e. node into which BatchNorm is
      nested.
    has_scaling: Whether the batch norm has scaling enabled.
    freeze_batch_norm_delay: How many steps to wait before freezing moving mean
      and variance and using them for batch normalization.
    is_training: Bool, true if training.

  Raises:
    ValueError: When operation type is not supported, or input and output tensor
      shapes mismatch for created operations: mul_fold, add_fold.

  Returns:
    A pair of Operations, the first is the original consumer node of the batch
      norm (../BatchNorm/batchnorm_1/add_1), the second is the consumer node of
      the folded graph (add_fold).
  """
  mul_scale_name = 'mul_1' if has_scaling else 'mul'
  mul_scale = graph.get_operation_by_name(context + 'BatchNorm/batchnorm_1/' +
                                          mul_scale_name)
  op_below = mul_scale.inputs[0].op
  # Skip over the BatchToSpace operation in the case of atrous convolutions.
  batch_to_space_op = None
  if op_below.type == 'BatchToSpaceND':
    batch_to_space_op = op_below
    op_below = op_below.inputs[0].op
  weights = op_below.inputs[1]
  match = _GetBatchNormParams(
      graph=graph, context=context, has_scaling=has_scaling)
  correction_scale, correction_recip, correction_offset = None, None, None
  if is_training:
    correction_scale, correction_recip, correction_offset = (
        _ComputeBatchNormCorrections(
            context=context,
            match=match,
            freeze_batch_norm_delay=freeze_batch_norm_delay))
  # Special handling for weights of depthwise convolution.
  if op_below.type == 'DepthwiseConv2dNative':
    new_shape = [
        weights.get_shape().as_list()[2],
        weights.get_shape().as_list()[3]
    ]
    scale_name = 'mul' if has_scaling else 'Rsqrt'
    scale = graph.get_operation_by_name(context + 'BatchNorm/batchnorm_1/' +
                                        scale_name)
    scale = array_ops.reshape(scale.outputs[0], new_shape,
                              context + 'scale_reshape')

    if correction_scale is not None:
      correction_scale = array_ops.reshape(correction_scale, new_shape,
                                           context + 'correction_reshape')
      with ops.device(mul_scale.device):
        weights = math_ops.multiply(correction_scale, weights,
                                    context + 'correction_mult')

    mul_fold = _CloneOp(mul_scale, context + 'mul_fold', [(0, weights),
                                                          (1, scale)])
  elif op_below.type in ['Conv2D', 'MatMul']:

    if correction_scale is not None:
      with ops.device(mul_scale.device):
        weights = math_ops.multiply(correction_scale, weights,
                                    context + 'correction_mult')
    mul_fold = _CloneOp(mul_scale, context + 'mul_fold', [(0, weights)])
  else:
    raise ValueError('Cannot handle operation of type: %s' % op_below.type)
  _AssertShapesMatch('mul_fold', mul_fold.inputs[0], mul_fold.outputs[0])

  conv_or_fc_folded = _CloneOp(op_below, op_below.name + '_Fold',
                               [(1, mul_fold.outputs[0])])

  add_shift = graph.get_operation_by_name(context +
                                          'BatchNorm/batchnorm_1/add_1')

  corrected_output = conv_or_fc_folded.outputs[0]
  # Copy the batch to space operation if we have a atrous convolution.
  if batch_to_space_op:
    corrected_output = array_ops.batch_to_space_nd(
        corrected_output,
        batch_to_space_op.inputs[1],
        batch_to_space_op.inputs[2],
        name=batch_to_space_op.name + '_Fold')
  if correction_offset is not None:
    with ops.device(conv_or_fc_folded.device):
      corrected_output = math_ops.multiply(correction_recip, corrected_output,
                                           context + 'post_conv_mul')
      corrected_output = math_ops.add(corrected_output, (correction_offset),
                                      context + 'correction_add')
  add_fold = _CloneOp(add_shift, context + 'add_fold', [(0, corrected_output)])
  _AssertShapesMatch('add_fold', add_fold.inputs[0], add_fold.outputs[0])
  return add_shift, add_fold