}
        }
        return std::nullopt;
      });
  if (!elem_type) return std::nullopt;
  return llvm::to_vector<8>(elem_type->getShape());
}

void ReplaceAllUsesWithCast(Value old_val, Value new_val) {
  if (old_val.use_empty()) return;
  auto cast_op =
      OpBuilder(old_val.getDefiningOp())
          .create<tensor::CastOp>(old_val.getLoc(), old_val.getType(), new_val);