if (!ranked_type)
    return context_op->emitOpError()
           << "A map_outside_compilation op's input and output types must be "
              "ranked tensors.";
  ArrayRef<int64_t> in_shape = ranked_type.getShape();
  if (in_shape.empty() || in_shape[0] < 0) {
    return context_op->emitOpError()
           << "A map_outside_compilation op's input and output shapes must "

    auto output_uniform_quantized_tensor_type = RankedTensorType::getChecked(
        rewriter.getUnknownLoc(),
        /*shape=*/
        mlir::cast<TensorType>(conv_output_value.getType()).getShape(),
        output_uniform_quantized_type);

    SmallVector<Type> new_conv_output_types = {
        output_uniform_quantized_tensor_type};
    auto new_conv_op_with_output_type =

  return static_cast<TF_DataType>(t->tensor->Type());
}

void TF_SetShape(TF_Tensor* t, const int64_t* dims, int num_dims) {
  tensorflow::down_cast<tensorflow::TensorInterface*>(t->tensor)->SetShape(
      dims, num_dims);
}

int TF_NumDims(const TF_Tensor* t) { return t->tensor->NumDims(); }

int64_t TF_Dim(const TF_Tensor* t, int dim_index) {
  return t->tensor->Dim(dim_index);
}

        }

        shape_inference::ShapeHandle handle;
        TF_RETURN_IF_ERROR(
            context->MakeShapeFromPartialTensorShape(arg_shape.shape, &handle));
        TF_RETURN_IF_ERROR(shape_refiner->SetShape(n, 0, handle));
      }
    }

    // Sometimes we have VariableShape nodes in while loop (after Enter nodes).
    // They won't be constant-folded because TensorFlow constant folding does

    return;
  }
  tensorflow::shape_inference::ShapeHandle new_shape =
      tensorflow::ShapeHandleFromDims(ic, num_dims, dims);
  status->status = graph->refiner.SetShape(node, output.index, new_shape);
}

int TF_GraphGetTensorNumDims(TF_Graph* graph, TF_Output output,
                             TF_Status* status) {
  Node* node = &output.oper->node;

  TF_Buffer* buf = TF_GetAllOpList();
  tensorflow::OpList op_list;
  EXPECT_TRUE(op_list.ParseFromArray(buf->data, buf->length));
  EXPECT_GT(op_list.op_size(), 0);
  TF_DeleteBuffer(buf);
}

TEST(CAPI, SetShape) {
  TF_Status* s = TF_NewStatus();
  TF_Graph* graph = TF_NewGraph();

  TF_Operation* feed = Placeholder(graph, s);
  ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
  TF_Output feed_out_0 = TF_Output{feed, 0};