class OpRegGen(transpiler.GenericTranspiler):
  """Transforms Python objects into TFR MLIR source code."""

  def transform_ast(self, node, ctx):
    gen = OpRegGenImpl(ctx)
    gen.visit(node)
    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