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 =