| 197 | |
| 198 | |
| 199 | class TensorProxy(StructInfoProxy): |
| 200 | shape: list[str | PrimExpr] | None |
| 201 | dtype: str |
| 202 | vdevice: str | None |
| 203 | ndim: int |
| 204 | |
| 205 | def __init__( |
| 206 | self, |
| 207 | shape: list[PrimExpr | str] | Expr | None = None, |
| 208 | dtype: str | None = None, |
| 209 | vdevice: str | None = None, |
| 210 | ndim: int = -1, |
| 211 | ) -> None: |
| 212 | if isinstance(shape, Expr): |
| 213 | if not isinstance(shape, ShapeExpr | Var): |
| 214 | raise ValueError( |
| 215 | "When the shape is an Expr, it must be a ShapeExpr or a Var with ShapeExpr " |
| 216 | f"value. But got: {shape} with type: {type(shape)}" |
| 217 | ) |
| 218 | if isinstance(shape, Var) and not isinstance(shape.struct_info, ShapeStructInfo): |
| 219 | raise ValueError( |
| 220 | "When the shape is a Var, it must have shape struct_info. But got " |
| 221 | f"{shape} with struct_info: {shape.struct_info}" |
| 222 | ) |
| 223 | self.shape = shape |
| 224 | self.dtype = dtype |
| 225 | self.vdevice = vdevice |
| 226 | self.ndim = ndim |
| 227 | |
| 228 | def get_symbolic_vars(self) -> set[str]: |
| 229 | if self.shape is None or isinstance(self.shape, Expr): |
| 230 | return {} |
| 231 | else: |
| 232 | return {s for s in self.shape if isinstance(s, str) and s.isidentifier()} |
| 233 | |
| 234 | def as_struct_info(self, dict_globals: dict[str, Any] | None = None) -> TensorStructInfo: |
| 235 | vdev = self.vdevice |
| 236 | if isinstance(self.vdevice, str): |
| 237 | if ":" in self.vdevice: |
| 238 | split_vdev = self.vdevice.split(":") |
| 239 | vdev = lookup_vdevice(split_vdev[0], int(split_vdev[1])) |
| 240 | else: |
| 241 | vdev = lookup_vdevice(self.vdevice, 0) |
| 242 | |
| 243 | if self.shape is None: |
| 244 | return TensorStructInfo(None, self.dtype, vdev, self.ndim) |
| 245 | elif isinstance(self.shape, ShapeExpr | Var): |
| 246 | return TensorStructInfo(self.shape, self.dtype, vdev, self.ndim) |
| 247 | else: |
| 248 | if dict_globals is None and any([isinstance(s, str) for s in self.shape]): |
| 249 | raise ValueError( |
| 250 | "String-defined shape expr is only allowed when parsing function parameters " |
| 251 | "and return annotations for TVMScript." |
| 252 | ) |
| 253 | shape = [_eval_shape(s, dict_globals) for s in self.shape] |
| 254 | return TensorStructInfo(shape, self.dtype, vdev, self.ndim) |
| 255 | |
| 256 | |