Note: tensor_metadata is only used in the case of a Tensor subclass, since fakifying a tensor subclass is not supported right now
(
graph_module: torch.fx.GraphModule,
val: Union[
Optional[FakeTensor], List[Optional[FakeTensor]], Tuple[Optional[FakeTensor]]
],
tensor_meta: Union[
Optional[TensorMetadata],
List[Optional[TensorMetadata]],
Tuple[Optional[TensorMetadata]],
],
)
| 267 | |
| 268 | |
| 269 | def make_alloc_node( |
| 270 | graph_module: torch.fx.GraphModule, |
| 271 | val: Union[ |
| 272 | Optional[FakeTensor], List[Optional[FakeTensor]], Tuple[Optional[FakeTensor]] |
| 273 | ], |
| 274 | tensor_meta: Union[ |
| 275 | Optional[TensorMetadata], |
| 276 | List[Optional[TensorMetadata]], |
| 277 | Tuple[Optional[TensorMetadata]], |
| 278 | ], |
| 279 | ) -> torch.fx.Node: |
| 280 | """ |
| 281 | Note: tensor_metadata is only used in the case of a Tensor subclass, since |
| 282 | fakifying a tensor subclass is not supported right now |
| 283 | """ |
| 284 | if val is None: |
| 285 | if tensor_meta is not None: |
| 286 | assert isinstance(tensor_meta, TensorMetadata) |
| 287 | alloc_spec = (tensor_meta.shape, tensor_meta.dtype) |
| 288 | else: |
| 289 | raise InternalError( |
| 290 | "Memory allocator node needs FakeTensor val or TensorMetadata to proceed" |
| 291 | ) |
| 292 | elif isinstance(val, FakeTensor): |
| 293 | alloc_spec = (val.shape, val.dtype) |
| 294 | else: |
| 295 | assert isinstance(val, list) or isinstance(val, tuple) |
| 296 | assert isinstance(tensor_meta, list) or isinstance(tensor_meta, tuple) |
| 297 | alloc_spec: List[memory.AllocSpec] = [] |
| 298 | for v, t in zip(val, tensor_meta): |
| 299 | if v is not None: |
| 300 | # pyre-fixme[6]: For 1st argument expected |
| 301 | # `Union[List[Tuple[List[int], dtype]], Tuple[List[int], dtype]]` but |
| 302 | # got `Tuple[Size, dtype]`. |
| 303 | alloc_spec.append((v.shape, v.dtype)) |
| 304 | elif t is not None: |
| 305 | # pyre-fixme[6]: For 1st argument expected |
| 306 | # `Union[List[Tuple[List[int], dtype]], Tuple[List[int], dtype]]` but |
| 307 | # got `Tuple[Size, dtype]`. |
| 308 | alloc_spec.append((t.shape, t.dtype)) |
| 309 | else: |
| 310 | raise InternalError( |
| 311 | "Memory allocator node needs FakeTensor val or TensorMetadata to proceed" |
| 312 | ) |
| 313 | |
| 314 | alloc = graph_module.graph.call_function(memory.alloc, (alloc_spec,)) |
| 315 | alloc.meta["val"] = val |
| 316 | alloc.meta["tensor_meta"] = tensor_meta |
| 317 | return alloc |
| 318 | |
| 319 | |
| 320 | class ToOutVarPass(PassBase): |
no test coverage detected