A helper function for capturing the shape as a tensor from a tensor value.
(x: torch.Tensor)
| 69 | |
| 70 | |
| 71 | def shape(x: torch.Tensor) -> Union[torch._C.Size, torch.Tensor]: |
| 72 | """ |
| 73 | A helper function for capturing the shape as a tensor from a tensor |
| 74 | value. |
| 75 | """ |
| 76 | tracer = DispatchTracer.get() |
| 77 | if tracer is None: |
| 78 | return x.shape |
| 79 | x = unwrap_functional(x) |
| 80 | if not isinstance(x, PythonTensor): |
| 81 | raise ExportError( |
| 82 | ExportErrorType.INVALID_INPUT_TYPE, |
| 83 | f"exir custom shape function only takes EXIR dispatch tensor, but got: {type(x)}", |
| 84 | ) |
| 85 | # TODO _shape_as_tensor should work with functional tensor but currently not. |
| 86 | # TODO torch.tensor() should succeed under functionalization but currently not. |
| 87 | # see: https://github.com/pytorch/pytorch/pull/76319 |
| 88 | tmp = torch.empty(len(x.shape), dtype=torch.int64) |
| 89 | for i, s in enumerate(x.shape): |
| 90 | tmp[i] = s |
| 91 | proxy = torch.ops.aten._shape_as_tensor.default(x.proxy) |
| 92 | return PythonTensor(unwrap_functional(tmp), proxy) |
| 93 | |
| 94 | |
| 95 | def _make_submodule( |
nothing calls this directly
no test coverage detected