REGISTER_OP("TestComplexTFOp")
    .Input("lhs: T")
    .Input("rhs: Tlen")
    .Output("output: N * T")
    .Attr("N: int >= 1")
    .Attr("T: numbertype")
    .Attr("Tlen: {int32, int64} = DT_INT64");

REGISTER_OP("TestNumAttrsOp")
    .Attr("x1: int = -10")
    .Attr("y1: int = 1")
    .Attr("x2: float = 0.0")
    .Attr("y2: float = -3.0");

REGISTER_OP("TestNonNumAttrsOp")
    .Attr("z: shape")
    .Attr("x: string = 'hello'")

  return z


@composite.Composite('TestTwoOutputsOp')
def _tfr_tensor_two_output(x):
  z = array_ops.Split(axis=0, value=x, num_split=2)
  return z[0], z[1]


@composite.Composite('TestNumAttrsOp')
def _tfr_tensor_tensor_with_cst(x1, y1, x2, y2):
  x = array_ops.OneHot(
      indices=[0, 2, -1, x1], depth=y1, on_value=True, off_value=False)
  (x, x2, y2)  # pylint: disable=pointless-statement
  return