auto input_type = operands[0].getType().dyn_cast<ShapedType>();
  if (!input_type || !input_type.hasRank()) {
    // If input is unranked, then so is output.
    inferredReturnTypes.assign(
        num_value, UnrankedTensorType::get(input_type.getElementType()));
    return success();
  }

  if (input_type.hasStaticShape() && input_type.getNumElements() <= 0) {

                            << n_init_values << ")";
  }

  auto input_ty_0 = inputs_ty[0].cast<ShapedType>();
  if (input_ty_0.hasStaticShape()) {
    for (int i = 1; i < n_inputs; ++i) {
      auto input_ty_i = inputs_ty[i].cast<ShapedType>();
      if (input_ty_i.hasStaticShape() &&
          input_ty_i.getShape() != input_ty_0.getShape()) {
        return op.emitOpError()

  llvm::SmallVector<int64_t, 4> input_shape(4, ShapedType::kDynamic);
  auto input_type = mlir::cast<TensorType>(op.getInput().getType());
  if (input_type.hasRank()) {
    if (input_type.getRank() != 4)
      return op.emitOpError()
             << "requires input to be a 4D tensor, but got " << input_type;

    int64_t input_batch = input_type.getDimSize(0);
    if (input_batch != ShapedType::kDynamic &&

      return failure();
    }

    const auto input_type = op.getLhs().getType().cast<TensorType>();
    if (!(input_type.getRank() == 2 || input_type.getRank() == 3)) {
      LLVM_DEBUG(llvm::dbgs() << "Input expected to have rank of 2 or 3. Got: "
                              << input_type << ".\n");
      return failure();
    }

    const Value filter = op.getRhs();

  @test_util.run_in_graph_and_eager_modes
  def test_qat_gather_and_conv_model(
      self,
  ):
    input_type = dtypes.int32
    model = self._create_simple_gather_and_conv_model(
        input_type,
        filter_shape=(2, 3, 3, 1024),
        is_qat_model=True,
    )

    saved_model_save.save(model, self._input_saved_model_path)

static bool AreInputDimensionsOneInAxes(Value input,
                                        const mlir::Attribute &axes) {
  RankedTensorType input_type =
      mlir::dyn_cast_or_null<RankedTensorType>(input.getType());
  if (!input_type) return false;
  auto type_shape = input_type.getShape();

  DenseIntElementsAttr axes_attr =
      mlir::dyn_cast_or_null<DenseIntElementsAttr>(axes);
  if (!axes_attr) return false;

    // If no padding is negative return the input as is.
    if (llvm::all_of(explicit_padding, [](int64_t pad) { return pad >= 0; })) {
      return value;
    }

    auto input_type = mlir::cast<RankedTensorType>(value.getType());
    auto input_shape = input_type.getShape();

    llvm::SmallVector<int64_t, 4> start;
    llvm::SmallVector<int64_t, 4> size;
    start.reserve(explicit_padding.size() / 2);

}

int TF_OperationNumInputs(TF_Operation* oper) {
  return oper->node.num_inputs();
}

TF_DataType TF_OperationInputType(TF_Input oper_in) {
  return static_cast<TF_DataType>(oper_in.oper->node.input_type(oper_in.index));
}

int TF_OperationInputListLength(TF_Operation* oper, const char* arg_name,
                                TF_Status* status) {
  NameRangeMap name_ranges;
  status->status =

    ArrayRef<int64_t> shape = arg_shapes[i];
    Type element_type;
    if (auto input_ty =
            mlir::dyn_cast<RankedTensorType>(func_type.getInput(i))) {
      if (input_ty.getRank() != shape.size()) {
        return failure();
      }
      element_type = input_ty.getElementType();
    } else {
      auto unranked_input_ty =

    auto* context = shape_refiner_->GetContext(&node);
    DataType dtype = node.input_type(idx);
    return ConvertDataTypeAndShape(dtype, context->input(idx),
                                   context->input_handle_shapes_and_types(idx),
                                   context, builder);
  }
  DataType dtype = node.properties()->input_types[idx];
  mlir::Type element_type;