from tensorflow.compiler.mlir.tfr.python.op_reg_gen import gen_register_op
from tensorflow.compiler.mlir.tfr.python.tfr_gen import tfr_gen_from_module
from tensorflow.python.platform import flags

FLAGS = flags.FLAGS

flags.DEFINE_string(
    'output', None,
    'Path to write the genereated register op file and MLIR file.')

flags.DEFINE_bool('gen_register_op', True,

def _override_test_attr_op():
  ret = array_ops.Const(value=100.0, dtype=dtypes.float32)
  return ret


def main(_):
  if FLAGS.gen_register_op:
    assert FLAGS.output.endswith('.cc')
    generated_code = gen_register_op(sys.modules[__name__], '_override_')
  else:
    assert FLAGS.output.endswith('.mlir')
    generated_code = tfr_gen_from_module(sys.modules[__name__], '_override_')

  dirname = os.path.dirname(FLAGS.output)

      strides=[1, stride_w, stride_h, 1],
      padding=padding,
      data_format='NHWC')


def main(_):
  if FLAGS.gen_register_op:
    assert FLAGS.output.endswith('.cc')
    generated_code = gen_register_op(sys.modules[__name__], '_composite_')
  else:
    assert FLAGS.output.endswith('.mlir')
    generated_code = tfr_gen_from_module(sys.modules[__name__], '_composite_',)

import sys

from tensorflow.compiler.mlir.python.mlir_wrapper import filecheck_wrapper as fw
from tensorflow.compiler.mlir.tfr.python import composite
from tensorflow.compiler.mlir.tfr.python.op_reg_gen import gen_register_op
from tensorflow.python.platform import test


Composite = composite.Composite


@composite.Composite(
    'TestNoOp', derived_attrs=['T: numbertype'], outputs=['o1: T'])
def _composite_no_op():
  pass

def _mirror_pad_grad_grad(op, grad):
  mode = op.get_attr('mode')
  return [gen_array_ops.mirror_pad(grad, op.inputs[1], mode=mode), None]


def main(_):
  if FLAGS.gen_register_op:
    assert FLAGS.output.endswith('.cc')
    generated_code = gen_register_op(sys.modules[__name__], '_composite_')
  else:
    assert FLAGS.output.endswith('.mlir')
    generated_code = tfr_gen_from_module(sys.modules[__name__], '_composite_')

    )

    registered_op = "registered_" + name
    native.genrule(
        name = registered_op,
        srcs = [],
        outs = [name + ".inc.cc"],
        cmd = "$(location %s) --output=$@ --gen_register_op=true" % gen_op_lib_exec,
        tools = [":" + gen_op_lib_exec],
        tags = tags,
    )

    native.cc_library(
        name = name + "_cc",
        testonly = test,
        srcs = [":" + registered_op],

    return gen.code_buffer


def op_reg_gen(func):
  """Parse a function and emit the TFR functions."""
  op_reg_code, _ = OpRegGen().transform(func, None)
  return op_reg_code


def gen_register_op(source, method_prefix=None):
  """Parse a python code and emit the TFR functions from a target class."""
  mlir_funcs = [
      op_reg_gen(func)
      for name, func in tf_inspect.getmembers(source, tf_inspect.isfunction)