Function DTensor

python/tvm/relax/script/parser/dist.py:69–103

(
    shape: list[PrimExpr | str] | None = None,
    dtype: str | None = None,
    device_mesh: DeviceMesh | str = DeviceMesh([], Range(0, 1)),
    placement: Placement | str = "",
    *,
    ndim: int = -1,
)

def DTensor(
    shape: list[PrimExpr | str] | None = None,
    dtype: str | None = None,
    device_mesh: DeviceMesh | str = DeviceMesh([], Range(0, 1)),
    placement: Placement | str = "",
    *,
    ndim: int = -1,
) -> DTensorProxy:
    # scalar tensor case
    if shape is not None and len(shape) == 0:
        shape = []
    if isinstance(shape, str) and dtype is None:
        dtype = shape
        shape = None

    if shape is not None and not isinstance(shape, tuple | list):
        raise ValueError(f"shape must be a list or tuple, but got: {shape}")
    if isinstance(device_mesh, str):
        if not IRBuilder.is_in_scope():
            return (
                DTensorProxy(
                    TensorProxy(shape, dtype, None, ndim), DeviceMesh([], Range(0, 1)), ""
                ),
            )
        name, index = device_mesh.split("[")
        index = int(index[:-1])
        frames = IRBuilder.current().frames
        for f in frames:
            if isinstance(f, IRModuleFrame):
                device_mesh = f.global_infos[name][index]
                break
        assert isinstance(device_mesh, DeviceMesh)
    if isinstance(placement, str):
        placement = Placement.from_text(placement)
    return DTensorProxy(TensorProxy(shape, dtype, None, ndim), device_mesh, placement)


__all__ = ["DTensor", "device_mesh"]