"@local_xla//xla/tsl/c:tsl_status_internal_headers",
    ],
    visibility = [
        "//tensorflow/python:__subpackages__",
    ],
)

cc_library(
    name = "tf_shape",
    srcs = ["tf_shape.cc"],
    hdrs = ["tf_shape.h"],
    copts = tf_copts(),
    visibility = ["//visibility:public"],
    deps = [
        ":c_api_macros",
        ":tf_shape_internal",
        "//tensorflow/core:framework",

    value {
      type: DT_STRING
    }
  }
  attr {
    key: "row_partition_types"
    value {
      list {
        s: ["ROW_SPLITS"]
      }
    }
  }
  attr {
    key: "Tshape"
    value {
      type: DT_INT64
    }
  }
  attr {
    key: "Tindex"
    value {
      type: DT_INT64
    }
  }
  attr {
    key: "T"
    value {
      type: DT_STRING

  TFE_Op* reshape_op = TFE_NewOp(tfe_context_, "Reshape", status_);
  CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_);
  TFE_OpSetAttrType(reshape_op, "T", TF_FLOAT);
  TFE_OpSetAttrType(reshape_op, "Tshape", TF_INT32);
  CheckOutputShapes(reshape_op,
                    /* input_shapes*/ {unknown_shape(), unknown_shape()},
                    /* input_tensors*/ {nullptr, tensor_1X6},

  TF_DeleteExecutionContext(ctx);
  return wrap(func);
}

TF_AbstractTensor* TF_AddFunctionParameter(TF_ExecutionContext* func,
                                           TF_DataType dtype, TF_Shape shape,
                                           TF_Status* s) {
  DCHECK_GE(shape.num_dims, -1);
  TracingTensorHandle* t;
  TracingContext* tracing_ctx = dyn_cast<TracingContext>(unwrap(func));
  if (!tracing_ctx) {

  %r1 = "tf.Reshape"(%in, %s1) {T = f32, Tshape = i32, device = ""} : (tensor<3x3x48x1xf32>, tensor<3xi32>) -> tensor<3x3x48xf32>
  %fq = "tf.FakeQuantWithMinMaxVarsPerChannel"(%r1, %mini, %maxi) {num_bits = 5, narrow_range = false} : (tensor<3x3x48xf32>, tensor<48xf32>, tensor<48xf32>) -> tensor<3x3x48xf32>
  %r2 = "tf.Reshape"(%fq, %s2) {T = f32, Tshape = i32, device = ""} : (tensor<3x3x48xf32>, tensor<4xi32>) -> tensor<3x3x48x1xf32>

  input: "BoxPredictor_4/ClassPredictor/act_quant/FakeQuantWithMinMaxVars"
  input: "BoxPredictor_4/Reshape_1/shape"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "Tshape"
    value {
      type: DT_INT32
    }
  }
}

# MLIR-LABEL: func @main

  input: "BoxPredictor_4/ClassPredictor/act_quant/FakeQuantWithMinMaxVars"
  input: "BoxPredictor_4/Reshape_1/shape"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "Tshape"
    value {
      type: DT_INT32
    }
  }
}

# MLIR-LABEL: func @main

  %r1 = "tf.Reshape"(%in, %s1) {T = f32, Tshape = i32, device = ""} : (tensor<3x3x48x1xf32>, tensor<3xi32>) -> tensor<3x3x48xf32>
  %fq = "tf.FakeQuantWithMinMaxVarsPerChannel"(%r1, %mini, %maxi) {num_bits = 3, narrow_range = false} : (tensor<3x3x48xf32>, tensor<48xf32>, tensor<48xf32>) -> tensor<3x3x48xf32>
  %r2 = "tf.Reshape"(%fq, %s2) {T = f32, Tshape = i32, device = ""} : (tensor<3x3x48xf32>, tensor<4xi32>) -> tensor<3x3x48x1xf32>

  input: "BoxPredictor_4/ClassPredictor/act_quant/FakeQuantWithMinMaxVars"
  input: "BoxPredictor_4/Reshape_1/shape"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "Tshape"
    value {
      type: DT_INT32
    }
  }
}

# MLIR-LABEL: func @main

  input: "BoxPredictor_4/ClassPredictor/act_quant/FakeQuantWithMinMaxVars"
  input: "BoxPredictor_4/Reshape_1/shape"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "Tshape"
    value {
      type: DT_INT32
    }
  }
}

# MLIR-LABEL: func @main