using ::tensorflow::calibrator::CalibrationStatisticsCollectorMinMax;
using ::tensorflow::calibrator::CalibrationStatisticsMap;

}  // namespace

REGISTER_OP("CalibrationStatisticsSaver")
    .Input("args: Tin")
    .Attr("Tin: list(type) >= 0")
    .Attr("ids: list(string) >= 1")
    .Attr("calibration_methods: list(int) >= 1")
    .Attr("output_file_path: string")
    .SetIsStateful()
    .Doc(R"doc(

  // CHECK: tfrt_fallback_async.batch_function device("/device:CPU:0") @batched_function
  // CHECK-SAME: Tin = [f32]
  // CHECK-SAME: Tout = [f32]
  // CHECK-SAME: allowed_batch_sizes = [6]
  // CHECK-SAME: batch_timeout_micros = 100000 : i64
  // CHECK-SAME: batching_queue = ""
  // CHECK-SAME: container = ""

    DefaultValuedAttr<BoolAttr, "false">:$used_by_host
  );

  let results = (outs
    TF_StrTensor:$array_key,
    TF_Tensor: $tensor_future
  );

  TF_DerivedOperandTypeListAttr Tin = TF_DerivedOperandTypeListAttr<0>;
  TF_DerivedResultTypeListAttr Tout = TF_DerivedResultTypeListAttr<0>;
}


def TF_IfrtCallOp : TF_Op<"IfrtCall", []> {
  let summary = "Invokes a program via IFRT on a device";

  let arguments = (ins
    Variadic<TensorOf<[TF_Float32, TF_Int64]>>:$inputs,

    StrAttr:$output_file_path,
    StrArrayAttr:$ids,
    I32ArrayAttr:$calibration_methods
  );

  TF_DerivedOperandTypeListAttr Tin = TF_DerivedOperandTypeListAttr<0>;
}

def TF_DumpTensorOp : TF_Op<"DumpTensor", []> {
  let summary = "Dump tensor proto.";

  let arguments = (ins
    Arg<TF_Tensor>:$input,

    StrAttr:$log_dir_path,

func.func @remove_unused_attr() {
  // CHECK: %out_op_chain = tfrt_fallback_async.executeop.seq(%arg0) key(0) cost({{.*}}) device("/device:CPU:0") "tf.SomeOp2"()
  "tf.SomeOp2"() {device = "/device:CPU:0", _output_shapes = ["tfshape$"], f.Tin = [f32], f._read_only_resource_inputs = []} : () -> ()
  func.return
}

// CHECK-LABEL: func @basic
func.func @basic(
    %arg0: tensor<3x1xf32>,