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


@Composite(
    'TestCompositeOp',
    inputs=['x: T', 'y: T'],
    attrs=['act: {"", "relu"}', 'trans: bool = true'],

==============================================================================*/

#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/op.h"

namespace tensorflow {

REGISTER_OP("TestNoOp");

REGISTER_OP("TestIdentityOp")
    .Input("input: T")
    .Output("output: T")
    .Attr("T: numbertype");

REGISTER_OP("TestIdentityNOp")
    .Input("input: N * T")
    .Output("output: N * T")

@Composite('TestInputNOp')
def _tfr_loc_test(x):
  n = 10
  x_sum = x[0]
  for i in range(1, n):
    x_sum = math_ops.Add(x_sum, x[i])
  return x_sum


#--- test fn for tfr tensors ---


@composite.Composite('TestNoOp')
def _tfr_tensor_empty_arg():
  pass


@composite.Composite('TestIdentityOp')
def _tfr_tensor_tensor(x):
  return x


@composite.Composite('TestIdentityNOp')
def _tfr_tensor_tensor_list(x):