| 40 | |
| 41 | |
| 42 | class DTensorProxy(StructInfoProxy): |
| 43 | tensor_sinfo_proxy: TensorProxy |
| 44 | device_mesh: DeviceMesh |
| 45 | placement: Placement |
| 46 | |
| 47 | def __init__( |
| 48 | self, |
| 49 | tensor_sinfo_proxy: TensorProxy, |
| 50 | device_mesh: DeviceMesh, |
| 51 | placement: Placement, |
| 52 | ) -> None: |
| 53 | self.device_mesh = device_mesh |
| 54 | self.placement = placement |
| 55 | self.tensor_sinfo_proxy = tensor_sinfo_proxy |
| 56 | super().__init__() |
| 57 | |
| 58 | def get_symbolic_vars(self) -> set[str]: |
| 59 | return self.tensor_sinfo_proxy.get_symbolic_vars() |
| 60 | |
| 61 | def as_struct_info(self, dict_globals: dict[str, Any] | None = None) -> TensorStructInfo: |
| 62 | return DTensorStructInfo( |
| 63 | self.tensor_sinfo_proxy.as_struct_info(dict_globals), |
| 64 | self.device_mesh, |
| 65 | self.placement, |
| 66 | ) |
| 67 | |
| 68 | |
| 69 | def DTensor( |
no outgoing calls
no test coverage detected
searching dependent graphs…