Type& shard_type) {
  RankedTensorType ranked_type = mlir::dyn_cast<RankedTensorType>(full_type);
  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) {

    // also populate the shape info to the op result. The shape of the
    // tensorlist is inferred from `num_elements` and `element_shape`.
    auto ranked_type = element_shape.getType().dyn_cast<RankedTensorType>();
    DenseIntElementsAttr dense_elem_attr;
    if ((ranked_type && ranked_type.getRank() == 0) ||
        !matchPattern(element_shape, m_Constant(&dense_elem_attr))) {
      // If no constant is spotted, just forward the operand.

// to expand at the given `axis`.
Type InferExpandDimsType(Type ty, int64_t axis, Builder *builder) {
  auto ranked_ty = mlir::dyn_cast<RankedTensorType>(ty);

  // Unranked type.
  if (!ranked_ty) return ty;

  auto shape = llvm::to_vector<4>(ranked_ty.getShape());
  if (axis < 0) axis += ranked_ty.getRank() + 1;

  shape.insert(shape.begin() + axis, 1);
  return tensorflow::GetTypeFromTFTensorShape(shape,