}
  *elem_type = RankedTensorType::get(elem_shape, buffer_type.getElementType());
  return success();
}

// Tries to infer the tensor array element shape.
std::optional<llvm::SmallVector<int64_t, 8>> GetTensorArrayElementShape(
    TF::TensorArrayV3Op ta, ModuleOp module) {
  auto element_shape = ta.getElementShapeAttr().cast<mlir::TF::ShapeAttr>();
  if (element_shape.hasStaticShape()) {
    auto shape = element_shape.getShape();