max_pool1 = gen_mnist_ops.new_max_pool(conv1, 2, 2, 2, 2, 'SAME')

    # output shape: [-1, 14, 14, 64]
    conv2 = gen_mnist_ops.new_conv2d(max_pool1, self.weights['f2'],
                                     self.biases['b2'], 1, 1, 1, 1, 'SAME',
                                     'RELU')

    # output shape: [-1, 7, 7, 64]
    max_pool2 = gen_mnist_ops.new_max_pool(conv2, 2, 2, 2, 2, 'SAME')

        'filter_width': 1,
        'filter_height': 1,
        'padding': 'SAME',
    }

    self._assertOpAndComposite([input_],
                               tf.function(gen_mnist_ops.new_max_pool),
                               ops_defs._composite_max_pool, kwargs)


if __name__ == '__main__':
  os.environ[
      'TF_MLIR_TFR_LIB_DIR'] = 'tensorflow/compiler/mlir/tfr/examples/mnist'

  b = math_ops.conj(op.inputs[1])
  grad_a = gen_math_ops.mat_mul(grad, b)
  grad_b = gen_math_ops.mat_mul(grad, a, transpose_a=True)
  return [grad_a, grad_b, bias_grad]


@Composite(
    'NewMaxPool',
    inputs=['input_: T'],
    attrs=[
        'stride_w: int', 'stride_h: int', 'filter_width: int',
        'filter_height: int', 'padding: {"SAME", "VALID"}'
    ],
    derived_attrs=['T: {float, int8}'],