Diff Coverage

Diff: origin/master...HEAD, staged and unstaged changes

Source File Diff Coverage (%) Missing Lines
hyper_parallel/platform/torch/dtensor.py 55.6% 24-26,30,246-247,249-250,269,274-275,289
hyper_parallel/platform/torch/dtensor.py
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

def _tensor_data_descriptor():
    """Return the raw Tensor.data descriptor from Tensor or its base classes."""
    for tensor_cls in getattr(Tensor, "__mro__", ()):  # tolerate tests patching Tensor with a mock object
        descriptor = vars(tensor_cls).get("data")
        if descriptor is not None:
            return descriptor
    descriptor = getattr(Tensor, "data", None)
    if descriptor is not None and hasattr(descriptor, "__get__"):
        return descriptor
    raise AttributeError("Tensor.data descriptor is not available")


class DTensorBase(Tensor):
    """torch dtensor base"""
242
243
244
245
246
247
248
249
250
251
252
253
254

    @property
    # pylint: disable=C2801
    def data(self):
        data_descriptor = _tensor_data_descriptor()
        local_data = data_descriptor.__get__(self._local_tensor, type(self._local_tensor))
        # pylint: disable=C0415
        from hyper_parallel.core.dtensor.dtensor import DTensor
        return DTensor.from_local(
            local_data,
            self._device_mesh,
            self._alias_placements(),
        )
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
        with getattr(torch, "_C").DisableTorchFunctionSubclass():
            data_descriptor.__set__(self, local_value)
            data_descriptor.__set__(self._local_tensor, local_value)
        if requires_grad is not None and hasattr(self, "requires_grad_"):
            self.requires_grad_(requires_grad)

    def cpu(self):
        """Move the local shard to CPU while preserving DTensor metadata."""
        # pylint: disable=C0415
        from hyper_parallel.core.dtensor.dtensor import DTensor
        return DTensor.from_local(
            self._local_tensor.cpu(),
            self._device_mesh,
            self._alias_placements(),
        )
285
286
287
288
289
290
291
292
293
        ``dtensor.data.cpu().numpy()``. NumPy has no distributed layout
        inference path; delegate to the local tensor to keep Tensor-like local
        data semantics.
        """
        return self._local_tensor.numpy(*args, **kwargs)

    @property
    def dtype(self) -> torch.dtype:
        """