        assert len(prof.profile.kernel_profiles[1].name) > 0
    except Exception as e:
        raise RuntimeError(
            "FusionDefinition did not run correctly with profile enabled! Error: "
            + str(e)
        )

def test_lru_fusion_cache_decorator():
    inputs = [
        torch.randn(2, 4, device="cuda"),
        torch.randn(2, 4, device="cuda"),
    ]

    def create_fd1(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.div(t0, t1)
        fd.add_output(t2)

    def create_fd2(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.div(t1, t0)
        fd.add_output(t2)

    @LruFusionCache(max_fusions=10)
    def create_fusion(create_fd: Callable):
        with FusionDefinition() as fd:
            create_fd(fd)
        return fd

    empty_stats = """Max Fusions Allowed: 10
The fusion cache is empty.\n"""
    assert create_fusion.stats() == empty_stats

    # Test LRU cache compilation
    for i in range(5):
        fd1 = create_fusion(create_fd1)
        outputs = fd1.execute(inputs)
        assert torch.allclose(outputs[0], inputs[0] / inputs[1])

        fd2 = create_fusion(create_fd2)
        outputs = fd2.execute(inputs)
        assert torch.allclose(outputs[0], inputs[1] / inputs[0])


def test_lru_fusion_cache_decorator():
    inputs = [
        torch.randn(2, 4, device="cuda"),
        torch.randn(2, 4, device="cuda"),
    ]

    def create_fd1(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.div(t0, t1)
        fd.add_output(t2)

    def create_fd2(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.div(t1, t0)
        fd.add_output(t2)

    @LruFusionCache(max_fusions=10)
    def create_fusion(create_fd: Callable):
        with FusionDefinition() as fd:
            create_fd(fd)
        return fd

    empty_stats = """Max Fusions Allowed: 10
The fusion cache is empty.\n"""
    assert create_fusion.stats() == empty_stats

    # Test LRU cache compilation
    for i in range(5):
        fd1 = create_fusion(create_fd1)
        outputs = fd1.execute(inputs)
        assert torch.allclose(outputs[0], inputs[0] / inputs[1])

        fd2 = create_fusion(create_fd2)
        outputs = fd2.execute(inputs)
        assert torch.allclose(outputs[0], inputs[1] / inputs[0])
        assert fd1.fusion == fd1.fusion
        assert fd2.fusion == fd2.fusion
        assert fd1.fusion != fd2.fusion

    expected_stats = """Max Fusions Allowed: 10
Total Fusions in Cache: 2


        assert fd1.fusion == fd1.fusion
        assert fd2.fusion == fd2.fusion
        assert fd1.fusion != fd2.fusion

    expected_stats = """Max Fusions Allowed: 10
Total Fusions in Cache: 2
Total Unique Fusions Compiled: 2
Cache Hits by LRU ordering:
\t0 -> 4 hits
\t1 -> 4 hits
Cache Lookups: 10
Cache Hits: 8
Hit Rate: 80%\n"""
    assert create_fusion.stats() == expected_stats

def test_lru_cache():
    bcast_input = torch.randn(1, 4, device="cuda")
    full_input = torch.randn(2, 4, device="cuda")

    def nvfuser_fusion_id0(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[1, -1],
            contiguity=[None, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T1 = fd.define_tensor(
            shape=[-1, -1],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T2 = fd.ops.sub(T0, T1)
        fd.add_output(T2)

    # nvfuser_fusion_id1 reverses the input argument order of nvfuser_fusion_id0
    def nvfuser_fusion_id1(fd: FusionDefinition) -> None:
        T1 = fd.define_tensor(
            shape=[-1, -1],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T0 = fd.define_tensor(
            shape=[1, -1],
            contiguity=[None, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T2 = fd.ops.sub(T0, T1)
        fd.add_output(T2)


    def nvfuser_fusion_id0(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[1, -1],
            contiguity=[None, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T1 = fd.define_tensor(
            shape=[-1, -1],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T2 = fd.ops.sub(T0, T1)
        fd.add_output(T2)

    # nvfuser_fusion_id1 reverses the input argument order of nvfuser_fusion_id0
    def nvfuser_fusion_id1(fd: FusionDefinition) -> None:
        T1 = fd.define_tensor(
            shape=[-1, -1],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T0 = fd.define_tensor(
            shape=[1, -1],
            contiguity=[None, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T2 = fd.ops.sub(T0, T1)
        fd.add_output(T2)

    cache = LRUCache(max_fusions=10)

    with FusionDefinition() as fd0:
        nvfuser_fusion_id0(fd0)

    assert cache.num_fusions() == 0
    fd0.fec = cache.cache_compile(fd0.fusion)
    del fd0._fusion
    # Check that LRUCache caches the first fusion
    assert cache.num_fusions() == 1

    with FusionDefinition() as fd1:
        nvfuser_fusion_id1(fd1)

    assert cache.num_fusions() == 1
    fd1.fec = cache.cache_compile(fd1.fusion)
    del fd1._fusion
    # Check that LRUCache creates a separate entry for second fusion
    assert cache.num_fusions() == 2


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest

import torch
import torch.distributed as dist

import transformer_engine.pytorch as te

from . import Parallelism
from .benchmark_utils import get_benchmark_fns
from enum import auto, Enum

compute_cap = torch.cuda.get_device_capability()

class ComputeType(Enum):
    FORWARD = auto()
    BACKWARD = auto()

# This benchmark is instrumented with cudaProfilerStart/Stop. Therefore, one
# can collect stats of the first few non-warmup benchmark iterations using
# ```bash
# mpirun -np <processes> nsys profile --capture-range=cudaProfilerApi --capture-range-end=repeat:<iterations> pytest tests/python/test_transformer_engine.py -k <filter> --only-mpi
# ```
# and then display the status using e.g. `nsys stats --report=cuda_gpu_kern_sum report1.nsys-rep`.
@pytest.mark.skipif(
    compute_cap >= (10, 0) and compute_cap < (12, 0),
    reason=f"TransformerEngine stalls this test on devices with compute capability {compute_cap}.",
)
@pytest.mark.mpi
@pytest.mark.parametrize(
    "compute_type",
    [ComputeType.FORWARD, ComputeType.BACKWARD],
    ids=lambda t: t.name,
)
@pytest.mark.parametrize(
    "parallelism",
    [Parallelism.TENSOR_PARALLEL, Parallelism.SEQUENCE_PARALLEL],
    ids=lambda p: p.name,
)
@pytest.mark.parametrize(
    "overlap",
    [
        pytest.param(False, id="nonoverlap"),
        pytest.param(
            True,
            id="overlap",
            marks=pytest.mark.skip(
                reason="`RuntimeError: /workspace/TransformerEngine/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers-host.cpp:321 in function create_communicator_grouped2: CUDA Error: invalid argument` in jit_python_distributed_tests_20_H100_TNVF"
            ),
        ),
    ],
)
def test_transformer_layer(
    setup_default_process_group,
    monkeypatch,
    benchmark,
    compute_type: ComputeType,
    parallelism: Parallelism,
    overlap: bool,
):
    if overlap and parallelism == Parallelism.TENSOR_PARALLEL:
        pytest.skip("Tensor parallelism doesn't support overlapping.")

    # Hyperparameters for GPT-3
    hidden_size = 12288
    num_heads = 96
    ffn_hidden_size = hidden_size * 4
    batch_size = 1
    sequence_length = 2048
    dtype = torch.bfloat16

    size = dist.get_world_size()

    transformer_layer = te.TransformerLayer(
        hidden_size,
        ffn_hidden_size,
        num_heads,
        # According to https://github.com/NVIDIA/TransformerEngine/issues/1350,
        # `attn_input_format` has to match the format of `transformer_layer`'s
        # input.
        attn_input_format="bshd",
        set_parallel_mode=True,
        sequence_parallel=(parallelism == Parallelism.SEQUENCE_PARALLEL),
        ub_tp_comm_overlap=overlap,
        tp_group=dist.group.WORLD,
    )
    transformer_layer.to(dtype).to("cuda")

    match parallelism:
        case Parallelism.TENSOR_PARALLEL:
            local_sequence_length = sequence_length
        case Parallelism.SEQUENCE_PARALLEL:
            assert sequence_length % size == 0
            local_sequence_length = sequence_length // size

    x = torch.randn(
        batch_size, local_sequence_length, hidden_size, dtype=dtype, device="cuda"
    )

    if overlap:
        # Similar to https://github.com/NVIDIA/TransformerEngine/blob/e7bfc0c547d63332e4f8d65e606dc69f4c22ffbe/examples/pytorch/comm_gemm_overlap/te_layer_with_overlap.py#L27-L29
        monkeypatch.setenv("CUDA_DEVICE_MAX_CONNECTIONS", "1")
        if not te.cpp_extensions.device_supports_multicast():
            monkeypatch.setenv("UB_SKIPMC", "1")

        te.module.base.initialize_ub(
            # Instructed by https://github.com/NVIDIA/TransformerEngine/blob/e7bfc0c547d63332e4f8d65e606dc69f4c22ffbe/transformer_engine/pytorch/module/base.py#L96-L99
            [batch_size * sequence_length, hidden_size],
            size,
            dtype=dtype,
            bootstrap_backend="nccl",
        )

    match compute_type:
        case ComputeType.FORWARD:
            warmup_fn, benchmark_fn = get_benchmark_fns(lambda: transformer_layer(x))
            y = warmup_fn()
            assert y.size() == torch.Size(
                [batch_size, local_sequence_length, hidden_size]
            )

            benchmark.pedantic(benchmark_fn, rounds=5)
        case ComputeType.BACKWARD:
            # Due to https://github.com/NVIDIA/TransformerEngine/issues/990, we can't repeatedly call
            # torch.autograd.backward to benchmark just backprop. As a
            # workaround, the code below runs forward before each backprop but
            # only measure the backprop time.
            def setup_fn():
                y = transformer_layer(x)
                dy = torch.rand_like(y)
                torch.cuda.synchronize()
                return (y, dy), {}

            warmup_fn, benchmark_fn = get_benchmark_fns(
                lambda y, dy: torch.autograd.backward(y, dy)
            )

            args, kwargs = setup_fn()
            warmup_fn(*args, **kwargs)

            benchmark.pedantic(
                benchmark_fn,
                setup=setup_fn,
                rounds=5,
            )

    if overlap:
        te.module.base.destroy_ub()


ComputeType

# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest

import torch
import torch.distributed as dist

import transformer_engine.pytorch as te

from . import Parallelism
from .benchmark_utils import get_benchmark_fns
from enum import auto, Enum

compute_cap = torch.cuda.get_device_capability()

class ComputeType(Enum):
    FORWARD = auto()
    BACKWARD = auto()

# This benchmark is instrumented with cudaProfilerStart/Stop. Therefore, one
# can collect stats of the first few non-warmup benchmark iterations using
# ```bash
# mpirun -np <processes> nsys profile --capture-range=cudaProfilerApi --capture-range-end=repeat:<iterations> pytest tests/python/test_transformer_engine.py -k <filter> --only-mpi
# ```
# and then display the status using e.g. `nsys stats --report=cuda_gpu_kern_sum report1.nsys-rep`.
@pytest.mark.skipif(
    compute_cap >= (10, 0) and compute_cap < (12, 0),
    reason=f"TransformerEngine stalls this test on devices with compute capability {compute_cap}.",
)
@pytest.mark.mpi
@pytest.mark.parametrize(
    "compute_type",
    [ComputeType.FORWARD, ComputeType.BACKWARD],
    ids=lambda t: t.name,
)
@pytest.mark.parametrize(
    "parallelism",


    ids=lambda t: t.name,
)
@pytest.mark.parametrize(
    "parallelism",
    [Parallelism.TENSOR_PARALLEL, Parallelism.SEQUENCE_PARALLEL],
    ids=lambda p: p.name,
)
@pytest.mark.parametrize(
    "overlap",
    [
        pytest.param(False, id="nonoverlap"),
        pytest.param(
            True,
            id="overlap",
            marks=pytest.mark.skip(
                reason="`RuntimeError: /workspace/TransformerEngine/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers-host.cpp:321 in function create_communicator_grouped2: CUDA Error: invalid argument` in jit_python_distributed_tests_20_H100_TNVF"
            ),
        ),
    ],
)
def test_transformer_layer(
    setup_default_process_group,
    monkeypatch,
    benchmark,
    compute_type: ComputeType,
    parallelism: Parallelism,
    overlap: bool,
):
    if overlap and parallelism == Parallelism.TENSOR_PARALLEL:
        pytest.skip("Tensor parallelism doesn't support overlapping.")

    # Hyperparameters for GPT-3
    hidden_size = 12288
    num_heads = 96
    ffn_hidden_size = hidden_size * 4
    batch_size = 1
    sequence_length = 2048
    dtype = torch.bfloat16

    size = dist.get_world_size()

    transformer_layer = te.TransformerLayer(
        hidden_size,
        ffn_hidden_size,
        num_heads,
        # According to https://github.com/NVIDIA/TransformerEngine/issues/1350,
        # `attn_input_format` has to match the format of `transformer_layer`'s
        # input.
        attn_input_format="bshd",
        set_parallel_mode=True,
        sequence_parallel=(parallelism == Parallelism.SEQUENCE_PARALLEL),
        ub_tp_comm_overlap=overlap,
        tp_group=dist.group.WORLD,
    )
    transformer_layer.to(dtype).to("cuda")

    match parallelism:
        case Parallelism.TENSOR_PARALLEL:
            local_sequence_length = sequence_length
        case Parallelism.SEQUENCE_PARALLEL:
            assert sequence_length % size == 0
            local_sequence_length = sequence_length // size

    x = torch.randn(
        batch_size, local_sequence_length, hidden_size, dtype=dtype, device="cuda"
    )

    if overlap:
        # Similar to https://github.com/NVIDIA/TransformerEngine/blob/e7bfc0c547d63332e4f8d65e606dc69f4c22ffbe/examples/pytorch/comm_gemm_overlap/te_layer_with_overlap.py#L27-L29
        monkeypatch.setenv("CUDA_DEVICE_MAX_CONNECTIONS", "1")
        if not te.cpp_extensions.device_supports_multicast():
            monkeypatch.setenv("UB_SKIPMC", "1")

        te.module.base.initialize_ub(
            # Instructed by https://github.com/NVIDIA/TransformerEngine/blob/e7bfc0c547d63332e4f8d65e606dc69f4c22ffbe/transformer_engine/pytorch/module/base.py#L96-L99
            [batch_size * sequence_length, hidden_size],
            size,
            dtype=dtype,
            bootstrap_backend="nccl",
        )

    match compute_type:
        case ComputeType.FORWARD:
            warmup_fn, benchmark_fn = get_benchmark_fns(lambda: transformer_layer(x))
            y = warmup_fn()
            assert y.size() == torch.Size(
                [batch_size, local_sequence_length, hidden_size]
            )

            benchmark.pedantic(benchmark_fn, rounds=5)
        case ComputeType.BACKWARD:
            # Due to https://github.com/NVIDIA/TransformerEngine/issues/990, we can't repeatedly call
            # torch.autograd.backward to benchmark just backprop. As a
            # workaround, the code below runs forward before each backprop but
            # only measure the backprop time.
            def setup_fn():
                y = transformer_layer(x)
                dy = torch.rand_like(y)
                torch.cuda.synchronize()
                return (y, dy), {}

            warmup_fn, benchmark_fn = get_benchmark_fns(
                lambda y, dy: torch.autograd.backward(y, dy)
            )

            args, kwargs = setup_fn()
            warmup_fn(*args, **kwargs)

            benchmark.pedantic(
                benchmark_fn,
                setup=setup_fn,
                rounds=5,
            )

    if overlap:
        te.module.base.destroy_ub()


            [batch_size * sequence_length, hidden_size],
            size,
            dtype=dtype,
            bootstrap_backend="nccl",
        )

    match compute_type:
        case ComputeType.FORWARD:
            warmup_fn, benchmark_fn = get_benchmark_fns(lambda: transformer_layer(x))
            y = warmup_fn()
            assert y.size() == torch.Size(
                [batch_size, local_sequence_length, hidden_size]
            )

            benchmark.pedantic(benchmark_fn, rounds=5)
        case ComputeType.BACKWARD:
            # Due to https://github.com/NVIDIA/TransformerEngine/issues/990, we can't repeatedly call
            # torch.autograd.backward to benchmark just backprop. As a
            # workaround, the code below runs forward before each backprop but
            # only measure the backprop time.
            def setup_fn():
                y = transformer_layer(x)
                dy = torch.rand_like(y)
                torch.cuda.synchronize()
                return (y, dy), {}

            warmup_fn, benchmark_fn = get_benchmark_fns(
                lambda y, dy: torch.autograd.backward(y, dy)
            )

            args, kwargs = setup_fn()
            warmup_fn(*args, **kwargs)

            benchmark.pedantic(
                benchmark_fn,
                setup=setup_fn,
                rounds=5,
            )

    if overlap:
        te.module.base.destroy_ub()


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import nvfuser_direct as nvfuser
from nvfuser_direct import FusionDefinition
from collections.abc import Iterable
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.placement_types import Placement, Shard
from typing import Callable, cast, TypeAlias

DTensorsKey: TypeAlias = tuple[tuple[str, str], ...]

def make_key_from_dtensors(dtensors: Iterable[DTensor]) -> DTensorsKey:
    key = tuple((repr(dt.device_mesh), repr(dt.placements)) for dt in dtensors)
    return key

class FusionDefinitionWrapper:
    def __init__(self, define_fusion: Callable):
        """Wraps a function that defines a fusion without `multidevice_schedule`."""
        # The complete FusionDefinition (w/ multidevice_scheduler) will have to
        # be created at "call" time according to the input DTensors.
        self._define_fusion = define_fusion

        # TODO: In theory, a FusionDefinitionWrapper can own multiple
        # `FusionDefinition`s, because different shardings lead to different
        # `multidevice_schedule`s. Currently, cache FusionDefinition based on input DTensors.
        self._fusion_definition_cache: dict[DTensorsKey, FusionDefinition] = {}

    def _create_fusion_definition(
        self, in_dtensors: Iterable[DTensor]
    ) -> FusionDefinition:
        with FusionDefinition() as fd:
            self._define_fusion(fd)
            self._multidevice_schedule(fd, in_dtensors)
        return fd

    def _multidevice_schedule(
        self, fd: FusionDefinition, in_dtensors: Iterable[DTensor]
    ) -> None:
        for in_tv, in_dtensor in zip(fd.fusion.inputs(), in_dtensors):
            mesh = nvfuser.multidevice.DeviceMesh(in_dtensor.device_mesh.mesh)

            in_tv.set_device_mesh(mesh)

            assert len(in_dtensor.placements) == 1, "Expect a 1D mesh"

            # Split and parallelize.
            # When the mesh is multi-dimensional, iterate through the
            # placements in descending order of Placement.dim.
            placement: Placement = in_dtensor.placements[0]
            if placement.is_shard():
                dim = cast(Shard, placement).dim
                in_tv.outer_split(dim, mesh.size)
                in_tv.axis(dim).parallelize(nvfuser.ParallelType.mesh_x)

    def _get_or_create_fusion_definition(
        self, in_dtensors: Iterable[DTensor]
    ) -> FusionDefinition:
        key = make_key_from_dtensors(in_dtensors)
        return self._fusion_definition_cache.setdefault(
            key, (lambda: self._create_fusion_definition(in_dtensors))()
        )

    def __call__(self, in_dtensors: Iterable[DTensor]) -> list[DTensor]:
        fusion_def = self._get_or_create_fusion_definition(in_dtensors)
        return nvfuser.execute_with_dtensors(fusion_def, in_dtensors)


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import nvfuser_direct as nvfuser
from nvfuser_direct import FusionDefinition
from collections.abc import Iterable
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.placement_types import Placement, Shard
from typing import Callable, cast, TypeAlias

DTensorsKey: TypeAlias = tuple[tuple[str, str], ...]

def make_key_from_dtensors(dtensors: Iterable[DTensor]) -> DTensorsKey:
    key = tuple((repr(dt.device_mesh), repr(dt.placements)) for dt in dtensors)
    return key

class FusionDefinitionWrapper:
    def __init__(self, define_fusion: Callable):
        """Wraps a function that defines a fusion without `multidevice_schedule`."""
        # The complete FusionDefinition (w/ multidevice_scheduler) will have to
        # be created at "call" time according to the input DTensors.
        self._define_fusion = define_fusion

        # TODO: In theory, a FusionDefinitionWrapper can own multiple
        # `FusionDefinition`s, because different shardings lead to different
        # `multidevice_schedule`s. Currently, cache FusionDefinition based on input DTensors.
        self._fusion_definition_cache: dict[DTensorsKey, FusionDefinition] = {}

    def _create_fusion_definition(
        self, in_dtensors: Iterable[DTensor]
    ) -> FusionDefinition:
        with FusionDefinition() as fd:
            self._define_fusion(fd)
            self._multidevice_schedule(fd, in_dtensors)


# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import nvfuser_direct as nvfuser
from nvfuser_direct import FusionDefinition
from collections.abc import Iterable
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.placement_types import Placement, Shard
from typing import Callable, cast, TypeAlias

DTensorsKey: TypeAlias = tuple[tuple[str, str], ...]

def make_key_from_dtensors(dtensors: Iterable[DTensor]) -> DTensorsKey:
    key = tuple((repr(dt.device_mesh), repr(dt.placements)) for dt in dtensors)
    return key

class FusionDefinitionWrapper:
    def __init__(self, define_fusion: Callable):
        """Wraps a function that defines a fusion without `multidevice_schedule`."""
        # The complete FusionDefinition (w/ multidevice_scheduler) will have to
        # be created at "call" time according to the input DTensors.
        self._define_fusion = define_fusion

        # TODO: In theory, a FusionDefinitionWrapper can own multiple
        # `FusionDefinition`s, because different shardings lead to different
        # `multidevice_schedule`s. Currently, cache FusionDefinition based on input DTensors.
        self._fusion_definition_cache: dict[DTensorsKey, FusionDefinition] = {}

    def _create_fusion_definition(
        self, in_dtensors: Iterable[DTensor]
    ) -> FusionDefinition:
        with FusionDefinition() as fd:
            self._define_fusion(fd)
            self._multidevice_schedule(fd, in_dtensors)
        return fd

    def _multidevice_schedule(
        self, fd: FusionDefinition, in_dtensors: Iterable[DTensor]
    ) -> None:
        for in_tv, in_dtensor in zip(fd.fusion.inputs(), in_dtensors):
            mesh = nvfuser.multidevice.DeviceMesh(in_dtensor.device_mesh.mesh)

            in_tv.set_device_mesh(mesh)

            assert len(in_dtensor.placements) == 1, "Expect a 1D mesh"

            # Split and parallelize.


DTensorsKey: TypeAlias = tuple[tuple[str, str], ...]

def make_key_from_dtensors(dtensors: Iterable[DTensor]) -> DTensorsKey:
    key = tuple((repr(dt.device_mesh), repr(dt.placements)) for dt in dtensors)
    return key

class FusionDefinitionWrapper:
    def __init__(self, define_fusion: Callable):
        """Wraps a function that defines a fusion without `multidevice_schedule`."""
        # The complete FusionDefinition (w/ multidevice_scheduler) will have to
        # be created at "call" time according to the input DTensors.
        self._define_fusion = define_fusion

        # TODO: In theory, a FusionDefinitionWrapper can own multiple
        # `FusionDefinition`s, because different shardings lead to different
        # `multidevice_schedule`s. Currently, cache FusionDefinition based on input DTensors.
        self._fusion_definition_cache: dict[DTensorsKey, FusionDefinition] = {}

    def _create_fusion_definition(
        self, in_dtensors: Iterable[DTensor]
    ) -> FusionDefinition:
        with FusionDefinition() as fd:
            self._define_fusion(fd)
            self._multidevice_schedule(fd, in_dtensors)
        return fd

    def _multidevice_schedule(
        self, fd: FusionDefinition, in_dtensors: Iterable[DTensor]
    ) -> None:
        for in_tv, in_dtensor in zip(fd.fusion.inputs(), in_dtensors):
            mesh = nvfuser.multidevice.DeviceMesh(in_dtensor.device_mesh.mesh)

            in_tv.set_device_mesh(mesh)

            assert len(in_dtensor.placements) == 1, "Expect a 1D mesh"

            # Split and parallelize.
            # When the mesh is multi-dimensional, iterate through the
            # placements in descending order of Placement.dim.
            placement: Placement = in_dtensor.placements[0]
            if placement.is_shard():
                dim = cast(Shard, placement).dim
                in_tv.outer_split(dim, mesh.size)
                in_tv.axis(dim).parallelize(nvfuser.ParallelType.mesh_x)


class FusionDefinitionWrapper:
    def __init__(self, define_fusion: Callable):
        """Wraps a function that defines a fusion without `multidevice_schedule`."""
        # The complete FusionDefinition (w/ multidevice_scheduler) will have to
        # be created at "call" time according to the input DTensors.
        self._define_fusion = define_fusion

        # TODO: In theory, a FusionDefinitionWrapper can own multiple
        # `FusionDefinition`s, because different shardings lead to different
        # `multidevice_schedule`s. Currently, cache FusionDefinition based on input DTensors.
        self._fusion_definition_cache: dict[DTensorsKey, FusionDefinition] = {}

    def _create_fusion_definition(
        self, in_dtensors: Iterable[DTensor]
    ) -> FusionDefinition:
        with FusionDefinition() as fd:
            self._define_fusion(fd)
            self._multidevice_schedule(fd, in_dtensors)
        return fd

    def _multidevice_schedule(
        self, fd: FusionDefinition, in_dtensors: Iterable[DTensor]
    ) -> None:
        for in_tv, in_dtensor in zip(fd.fusion.inputs(), in_dtensors):
            mesh = nvfuser.multidevice.DeviceMesh(in_dtensor.device_mesh.mesh)

            in_tv.set_device_mesh(mesh)

            assert len(in_dtensor.placements) == 1, "Expect a 1D mesh"

            # Split and parallelize.
            # When the mesh is multi-dimensional, iterate through the
            # placements in descending order of Placement.dim.
            placement: Placement = in_dtensor.placements[0]
            if placement.is_shard():
                dim = cast(Shard, placement).dim
                in_tv.outer_split(dim, mesh.size)
                in_tv.axis(dim).parallelize(nvfuser.ParallelType.mesh_x)

    def _get_or_create_fusion_definition(
        self, in_dtensors: Iterable[DTensor]
    ) -> FusionDefinition:
        key = make_key_from_dtensors(in_dtensors)
        return self._fusion_definition_cache.setdefault(
            key, (lambda: self._create_fusion_definition(in_dtensors))()
        )

    def __call__(self, in_dtensors: Iterable[DTensor]) -> list[DTensor]:
        fusion_def = self._get_or_create_fusion_definition(in_dtensors)
        return nvfuser.execute_with_dtensors(fusion_def, in_dtensors)


    def _multidevice_schedule(
        self, fd: FusionDefinition, in_dtensors: Iterable[DTensor]
    ) -> None:
        for in_tv, in_dtensor in zip(fd.fusion.inputs(), in_dtensors):
            mesh = nvfuser.multidevice.DeviceMesh(in_dtensor.device_mesh.mesh)

            in_tv.set_device_mesh(mesh)

            assert len(in_dtensor.placements) == 1, "Expect a 1D mesh"

            # Split and parallelize.
            # When the mesh is multi-dimensional, iterate through the
            # placements in descending order of Placement.dim.
            placement: Placement = in_dtensor.placements[0]
            if placement.is_shard():
                dim = cast(Shard, placement).dim
                in_tv.outer_split(dim, mesh.size)
                in_tv.axis(dim).parallelize(nvfuser.ParallelType.mesh_x)

    def _get_or_create_fusion_definition(
        self, in_dtensors: Iterable[DTensor]
    ) -> FusionDefinition:
        key = make_key_from_dtensors(in_dtensors)
        return self._fusion_definition_cache.setdefault(
            key, (lambda: self._create_fusion_definition(in_dtensors))()
        )

    def __call__(self, in_dtensors: Iterable[DTensor]) -> list[DTensor]:
        fusion_def = self._get_or_create_fusion_definition(in_dtensors)
        return nvfuser.execute_with_dtensors(fusion_def, in_dtensors)


            assert len(in_dtensor.placements) == 1, "Expect a 1D mesh"

            # Split and parallelize.
            # When the mesh is multi-dimensional, iterate through the
            # placements in descending order of Placement.dim.
            placement: Placement = in_dtensor.placements[0]
            if placement.is_shard():
                dim = cast(Shard, placement).dim
                in_tv.outer_split(dim, mesh.size)
                in_tv.axis(dim).parallelize(nvfuser.ParallelType.mesh_x)

    def _get_or_create_fusion_definition(
        self, in_dtensors: Iterable[DTensor]
    ) -> FusionDefinition:
        key = make_key_from_dtensors(in_dtensors)
        return self._fusion_definition_cache.setdefault(
            key, (lambda: self._create_fusion_definition(in_dtensors))()
        )

    def __call__(self, in_dtensors: Iterable[DTensor]) -> list[DTensor]:
        fusion_def = self._get_or_create_fusion_definition(in_dtensors)
        return nvfuser.execute_with_dtensors(fusion_def, in_dtensors)


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest

import torch
import torch.distributed as dist

from torch.distributed.tensor import DTensor
from torch.distributed.tensor.placement_types import Shard, Replicate
from .fusion_definition_wrapper import FusionDefinitionWrapper
from .linear import LinearFunction
from nvfuser_direct import FusionDefinition, DataType

@pytest.mark.mpi
def test_plus_one(setup_default_process_group, multidevice_test):
    def define_fusion(fd: FusionDefinition):
        inp = fd.define_tensor((-1, -1), contiguity=False, dtype=DataType.Float)
        one = fd.define_scalar(1.0, dtype=DataType.Float)
        out = fd.ops.add(inp, one)
        fd.add_output(out)

    op = FusionDefinitionWrapper(define_fusion)

    num_devices = dist.get_world_size()

    in_tensor = torch.randn(num_devices, 4)
    mesh = dist.device_mesh.init_device_mesh(
        "cuda", [num_devices], mesh_dim_names=["tp"]
    )
    in_dtensor = dist.tensor.distribute_tensor(in_tensor, mesh, [Shard(0)])

    (out_dtensor,) = op([in_dtensor])
    torch.testing.assert_close(out_dtensor.to_local(), in_dtensor.to_local() + 1)
    assert out_dtensor.device_mesh == in_dtensor.device_mesh
    assert out_dtensor.placements == in_dtensor.placements
    assert (
        out_dtensor.device_mesh.mesh_dim_names == in_dtensor.device_mesh.mesh_dim_names
    )

@pytest.mark.mpi
def test_column_parallel_linear(setup_default_process_group, multidevice_test):
    d, b, s, e = dist.get_world_size(), 2, 3, 5

    mesh = dist.device_mesh.init_device_mesh("cuda", [d])

    inp_tensor = torch.randint(
        -4, 4, (b, s, e), dtype=torch.bfloat16, requires_grad=True
    )
    weight_tensor = torch.randint(
        -4, 4, (d * e, e), dtype=torch.bfloat16, requires_grad=True
    )

    inp_dtensor = dist.tensor.distribute_tensor(inp_tensor, mesh, [Replicate()])
    weight_dtensor = dist.tensor.distribute_tensor(weight_tensor, mesh, [Shard(0)])

    def assert_close(expected_tensor: torch.Tensor, dtensor: DTensor):
        torch.testing.assert_close(expected_tensor, dtensor.to_local().cpu())

    out_tensor = torch.nn.functional.linear(inp_tensor, weight_tensor)
    out_dtensor = LinearFunction.apply(inp_dtensor, weight_dtensor)
    rank = dist.get_rank()
    assert_close(out_tensor.split(e, dim=-1)[rank], out_dtensor)

    (expected_grad_x, expected_grad_w) = torch.autograd.grad(
        out_tensor,
        (inp_tensor, weight_tensor),
        torch.ones_like(out_tensor),
    )
    (grad_x, grad_w) = torch.autograd.grad(
        out_dtensor,
        (inp_dtensor, weight_dtensor),
        torch.ones_like(out_dtensor),
    )
    assert_close(expected_grad_x, grad_x)
    assert_close(expected_grad_w.split(e, dim=0)[rank], grad_w)

@pytest.mark.mpi
def test_row_parallel_linear(setup_default_process_group, multidevice_test):
    d, b, s, e = dist.get_world_size(), 2, 3, 5

    mesh = dist.device_mesh.init_device_mesh("cuda", [d])

    inp_tensor = torch.randint(
        -4, 4, (b, s, d * e), dtype=torch.bfloat16, requires_grad=True
    )
    weight_tensor = torch.randint(
        -4, 4, (e, d * e), dtype=torch.bfloat16, requires_grad=True
    )

    inp_dtensor = dist.tensor.distribute_tensor(inp_tensor, mesh, [Shard(-1)])
    weight_dtensor = dist.tensor.distribute_tensor(weight_tensor, mesh, [Shard(-1)])

    def assert_close(expected_tensor: torch.Tensor, dtensor: DTensor):
        torch.testing.assert_close(expected_tensor, dtensor.to_local().cpu())

    out_tensor = torch.nn.functional.linear(inp_tensor, weight_tensor)
    out_dtensor = LinearFunction.apply(inp_dtensor, weight_dtensor)
    assert_close(out_tensor, out_dtensor)

    (expected_grad_x, expected_grad_w) = torch.autograd.grad(
        out_tensor,
        (inp_tensor, weight_tensor),
        torch.ones_like(out_tensor),
    )
    (grad_x, grad_w) = torch.autograd.grad(
        out_dtensor,
        (inp_dtensor, weight_dtensor),
        torch.ones_like(out_dtensor),
    )
    rank = dist.get_rank()
    assert_close(expected_grad_x.split(e, dim=-1)[rank], grad_x)
    assert_close(expected_grad_w.split(e, dim=-1)[rank], grad_w)


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest

import torch
import torch.distributed as dist

from torch.distributed.tensor import DTensor
from torch.distributed.tensor.placement_types import Shard, Replicate
from .fusion_definition_wrapper import FusionDefinitionWrapper
from .linear import LinearFunction
from nvfuser_direct import FusionDefinition, DataType

@pytest.mark.mpi
def test_plus_one(setup_default_process_group, multidevice_test):
    def define_fusion(fd: FusionDefinition):
        inp = fd.define_tensor((-1, -1), contiguity=False, dtype=DataType.Float)
        one = fd.define_scalar(1.0, dtype=DataType.Float)
        out = fd.ops.add(inp, one)
        fd.add_output(out)

    op = FusionDefinitionWrapper(define_fusion)

    num_devices = dist.get_world_size()

    in_tensor = torch.randn(num_devices, 4)
    mesh = dist.device_mesh.init_device_mesh(
        "cuda", [num_devices], mesh_dim_names=["tp"]
    )
    in_dtensor = dist.tensor.distribute_tensor(in_tensor, mesh, [Shard(0)])

    (out_dtensor,) = op([in_dtensor])
    torch.testing.assert_close(out_dtensor.to_local(), in_dtensor.to_local() + 1)
    assert out_dtensor.device_mesh == in_dtensor.device_mesh
    assert out_dtensor.placements == in_dtensor.placements
    assert (
        out_dtensor.device_mesh.mesh_dim_names == in_dtensor.device_mesh.mesh_dim_names
    )

@pytest.mark.mpi
def test_column_parallel_linear(setup_default_process_group, multidevice_test):
    d, b, s, e = dist.get_world_size(), 2, 3, 5

    mesh = dist.device_mesh.init_device_mesh("cuda", [d])

    inp_tensor = torch.randint(
        -4, 4, (b, s, e), dtype=torch.bfloat16, requires_grad=True
    )
    weight_tensor = torch.randint(
        -4, 4, (d * e, e), dtype=torch.bfloat16, requires_grad=True
    )

    inp_dtensor = dist.tensor.distribute_tensor(inp_tensor, mesh, [Replicate()])
    weight_dtensor = dist.tensor.distribute_tensor(weight_tensor, mesh, [Shard(0)])

    def assert_close(expected_tensor: torch.Tensor, dtensor: DTensor):
        torch.testing.assert_close(expected_tensor, dtensor.to_local().cpu())


    op = FusionDefinitionWrapper(define_fusion)

    num_devices = dist.get_world_size()

    in_tensor = torch.randn(num_devices, 4)
    mesh = dist.device_mesh.init_device_mesh(
        "cuda", [num_devices], mesh_dim_names=["tp"]
    )
    in_dtensor = dist.tensor.distribute_tensor(in_tensor, mesh, [Shard(0)])

    (out_dtensor,) = op([in_dtensor])
    torch.testing.assert_close(out_dtensor.to_local(), in_dtensor.to_local() + 1)
    assert out_dtensor.device_mesh == in_dtensor.device_mesh
    assert out_dtensor.placements == in_dtensor.placements
    assert (
        out_dtensor.device_mesh.mesh_dim_names == in_dtensor.device_mesh.mesh_dim_names
    )

@pytest.mark.mpi
def test_column_parallel_linear(setup_default_process_group, multidevice_test):
    d, b, s, e = dist.get_world_size(), 2, 3, 5

    mesh = dist.device_mesh.init_device_mesh("cuda", [d])

    inp_tensor = torch.randint(
        -4, 4, (b, s, e), dtype=torch.bfloat16, requires_grad=True
    )
    weight_tensor = torch.randint(
        -4, 4, (d * e, e), dtype=torch.bfloat16, requires_grad=True
    )

    inp_dtensor = dist.tensor.distribute_tensor(inp_tensor, mesh, [Replicate()])
    weight_dtensor = dist.tensor.distribute_tensor(weight_tensor, mesh, [Shard(0)])

    def assert_close(expected_tensor: torch.Tensor, dtensor: DTensor):
        torch.testing.assert_close(expected_tensor, dtensor.to_local().cpu())

    out_tensor = torch.nn.functional.linear(inp_tensor, weight_tensor)
    out_dtensor = LinearFunction.apply(inp_dtensor, weight_dtensor)
    rank = dist.get_rank()
    assert_close(out_tensor.split(e, dim=-1)[rank], out_dtensor)

    (expected_grad_x, expected_grad_w) = torch.autograd.grad(
        out_tensor,
        (inp_tensor, weight_tensor),
        torch.ones_like(out_tensor),
    )
    (grad_x, grad_w) = torch.autograd.grad(
        out_dtensor,
        (inp_dtensor, weight_dtensor),
        torch.ones_like(out_dtensor),
    )
    assert_close(expected_grad_x, grad_x)
    assert_close(expected_grad_w.split(e, dim=0)[rank], grad_w)

@pytest.mark.mpi
def test_row_parallel_linear(setup_default_process_group, multidevice_test):
    d, b, s, e = dist.get_world_size(), 2, 3, 5

    mesh = dist.device_mesh.init_device_mesh("cuda", [d])

    inp_tensor = torch.randint(
        -4, 4, (b, s, d * e), dtype=torch.bfloat16, requires_grad=True
    )
    weight_tensor = torch.randint(
        -4, 4, (e, d * e), dtype=torch.bfloat16, requires_grad=True
    )

    inp_dtensor = dist.tensor.distribute_tensor(inp_tensor, mesh, [Shard(-1)])
    weight_dtensor = dist.tensor.distribute_tensor(weight_tensor, mesh, [Shard(-1)])

    def assert_close(expected_tensor: torch.Tensor, dtensor: DTensor):
        torch.testing.assert_close(expected_tensor, dtensor.to_local().cpu())


    out_tensor = torch.nn.functional.linear(inp_tensor, weight_tensor)
    out_dtensor = LinearFunction.apply(inp_dtensor, weight_dtensor)
    rank = dist.get_rank()
    assert_close(out_tensor.split(e, dim=-1)[rank], out_dtensor)

    (expected_grad_x, expected_grad_w) = torch.autograd.grad(
        out_tensor,
        (inp_tensor, weight_tensor),
        torch.ones_like(out_tensor),
    )
    (grad_x, grad_w) = torch.autograd.grad(
        out_dtensor,
        (inp_dtensor, weight_dtensor),
        torch.ones_like(out_dtensor),
    )
    assert_close(expected_grad_x, grad_x)
    assert_close(expected_grad_w.split(e, dim=0)[rank], grad_w)

@pytest.mark.mpi
def test_row_parallel_linear(setup_default_process_group, multidevice_test):
    d, b, s, e = dist.get_world_size(), 2, 3, 5

    mesh = dist.device_mesh.init_device_mesh("cuda", [d])

    inp_tensor = torch.randint(
        -4, 4, (b, s, d * e), dtype=torch.bfloat16, requires_grad=True
    )
    weight_tensor = torch.randint(
        -4, 4, (e, d * e), dtype=torch.bfloat16, requires_grad=True
    )

    inp_dtensor = dist.tensor.distribute_tensor(inp_tensor, mesh, [Shard(-1)])
    weight_dtensor = dist.tensor.distribute_tensor(weight_tensor, mesh, [Shard(-1)])

    def assert_close(expected_tensor: torch.Tensor, dtensor: DTensor):
        torch.testing.assert_close(expected_tensor, dtensor.to_local().cpu())

    out_tensor = torch.nn.functional.linear(inp_tensor, weight_tensor)
    out_dtensor = LinearFunction.apply(inp_dtensor, weight_dtensor)
    assert_close(out_tensor, out_dtensor)

    (expected_grad_x, expected_grad_w) = torch.autograd.grad(
        out_tensor,
        (inp_tensor, weight_tensor),
        torch.ones_like(out_tensor),
    )
    (grad_x, grad_w) = torch.autograd.grad(
        out_dtensor,
        (inp_dtensor, weight_dtensor),
        torch.ones_like(out_dtensor),
    )
    rank = dist.get_rank()
    assert_close(expected_grad_x.split(e, dim=-1)[rank], grad_x)
    assert_close(expected_grad_w.split(e, dim=-1)[rank], grad_w)


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest

import torch
import torch.distributed as dist

from torch.distributed.tensor import DTensor
from torch.distributed.tensor.placement_types import Shard, Replicate
from .fusion_definition_wrapper import FusionDefinitionWrapper
from .linear import LinearFunction
from nvfuser_direct import FusionDefinition, DataType

@pytest.mark.mpi
def test_plus_one(setup_default_process_group, multidevice_test):
    def define_fusion(fd: FusionDefinition):
        inp = fd.define_tensor((-1, -1), contiguity=False, dtype=DataType.Float)
        one = fd.define_scalar(1.0, dtype=DataType.Float)
        out = fd.ops.add(inp, one)
        fd.add_output(out)

    op = FusionDefinitionWrapper(define_fusion)

    num_devices = dist.get_world_size()

    in_tensor = torch.randn(num_devices, 4)
    mesh = dist.device_mesh.init_device_mesh(
        "cuda", [num_devices], mesh_dim_names=["tp"]
    )
    in_dtensor = dist.tensor.distribute_tensor(in_tensor, mesh, [Shard(0)])

    (out_dtensor,) = op([in_dtensor])
    torch.testing.assert_close(out_dtensor.to_local(), in_dtensor.to_local() + 1)
    assert out_dtensor.device_mesh == in_dtensor.device_mesh
    assert out_dtensor.placements == in_dtensor.placements
    assert (
        out_dtensor.device_mesh.mesh_dim_names == in_dtensor.device_mesh.mesh_dim_names
    )


    assert_close(expected_grad_x, grad_x)
    assert_close(expected_grad_w.split(e, dim=0)[rank], grad_w)

@pytest.mark.mpi
def test_row_parallel_linear(setup_default_process_group, multidevice_test):
    d, b, s, e = dist.get_world_size(), 2, 3, 5

    mesh = dist.device_mesh.init_device_mesh("cuda", [d])

    inp_tensor = torch.randint(
        -4, 4, (b, s, d * e), dtype=torch.bfloat16, requires_grad=True
    )
    weight_tensor = torch.randint(
        -4, 4, (e, d * e), dtype=torch.bfloat16, requires_grad=True
    )

    inp_dtensor = dist.tensor.distribute_tensor(inp_tensor, mesh, [Shard(-1)])
    weight_dtensor = dist.tensor.distribute_tensor(weight_tensor, mesh, [Shard(-1)])

    def assert_close(expected_tensor: torch.Tensor, dtensor: DTensor):
        torch.testing.assert_close(expected_tensor, dtensor.to_local().cpu())

    out_tensor = torch.nn.functional.linear(inp_tensor, weight_tensor)
    out_dtensor = LinearFunction.apply(inp_dtensor, weight_dtensor)
    assert_close(out_tensor, out_dtensor)

    (expected_grad_x, expected_grad_w) = torch.autograd.grad(
        out_tensor,
        (inp_tensor, weight_tensor),
        torch.ones_like(out_tensor),
    )
    (grad_x, grad_w) = torch.autograd.grad(
        out_dtensor,
        (inp_dtensor, weight_dtensor),
        torch.ones_like(out_dtensor),
    )
    rank = dist.get_rank()
    assert_close(expected_grad_x.split(e, dim=-1)[rank], grad_x)
    assert_close(expected_grad_w.split(e, dim=-1)[rank], grad_w)


tests.python.multidevice.conftest

# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

# Run command:
# mpirun -np [num_devices] pytest tests/python/multidevice/[test_name].py --only-mpi -s

import os
import pytest
from typing import Iterable

import torch
import torch.distributed as dist

import nvfuser_direct as nvfuser

class MultideviceTest:
    def __init__(self):
        os.environ["NVIDIA_TF32_OVERRIDE"] = "0"

        self._communicator = nvfuser.multidevice.Communicator.instance()

        torch.cuda.set_device(self._communicator.local_rank())

        # This way, when individual tests create unsharded input, each rank
        # receives the same data.
        torch.manual_seed(0)

    @property
    def communicator(self):
        return self._communicator

    @property
    def size(self):
        return self._communicator.size()

    @property
    def rank(self):
        return self._communicator.rank()

    @property
    def local_size(self):
        return self._communicator.local_size()

    @property
    def local_rank(self):
        return self._communicator.local_rank()

    def shard_tensor_1d(self, t: torch.Tensor, dim: int, mesh) -> torch.Tensor:
        """Shard tensor along a single dimension (1D sharding only).

        Args:
            t: Tensor to shard (preferably on CPU for memory efficiency)
            dim: Dimension to shard along
            mesh: DeviceMesh to use for sharding

        Returns:
            Sharded tensor on current GPU device

        Example:
            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(num_devices))
            unsharded = torch.randn(num_devices, 4)
            sharded = self.shard_tensor_1d(unsharded, 0, mesh)
        """
        assert t.is_cpu, (
            "This is not strictly required but it's a general good practice "
            "for unit tests to create unsharded data on CPU to reduce GPU "
            "memory footprint."
        )
        return mesh.shard_tensor(t, dim).cuda(self.local_rank)

    def shard_tensor(self, t: torch.Tensor, tv) -> torch.Tensor:
        """Shard tensor using TensorView's parallelization and device mesh.

        Args:
            t: Tensor to shard (preferably on CPU for memory efficiency)
            tv: TensorView with device mesh and parallelization information

        Returns:
            Sharded tensor on current GPU device

        Example:
            with nvfuser.FusionDefinition() as fd:
                inp_tv = fd.define_tensor([-1, -1])
                inp_tv.set_device_mesh(mesh)
                inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)
                # ... rest of fusion

            unsharded = torch.randn(num_devices, 4)
            sharded = self.shard_tensor(unsharded, inp_tv)
        """
        assert t.is_cpu, (
            "This is not strictly required but it's a general good practice "
            "for unit tests to create unsharded data on CPU to reduce GPU "
            "memory footprint."
        )

        sharded = nvfuser.multidevice.shard_tensor(t, tv)
        return sharded.cuda(self.local_rank)

@pytest.fixture
def multidevice_test():
    fixture = MultideviceTest()
    yield fixture
    fixture.communicator.barrier()

def get_env(envs: Iterable[str], /, *, default: str) -> str:
    for env in envs:
        if value := os.environ.get(env):
            return value
    return default

# Set up the default process group for torch APIs like
# dist.device_mesh.init_device_mesh.
#
# This fixture is used by multi-GPU tests that use torch.distributed directly.
#
# I use "session" instead of "module" because
# https://github.com/pytorch/pytorch/issues/119196 reported race conditions
# when reinitializing process groups.
@pytest.fixture(scope="session")
def setup_default_process_group():
    # I avoided using nvfuser.Communicator to minimize fixture dependencies on
    # nvFuser. This makes the transition from legacy bindings to direct
    # bindings easier.
    rank = int(get_env(["OMPI_COMM_WORLD_RANK", "RANK"], default="0"))
    world_size = int(get_env(["OMPI_COMM_WORLD_SIZE", "WORLD_SIZE"], default="1"))
    local_rank = int(get_env(["OMPI_COMM_WORLD_LOCAL_RANK", "LOCAL_RANK"], default="0"))

    torch.cuda.set_device(local_rank)

    # The default port as used by https://github.com/pytorch/pytorch/blob/45a8b5682eb69d865cbf68c7f2f689b56b4efd53/torch/csrc/distributed/c10d/TCPStore.hpp#L51.
    dist.init_process_group(
        backend="nccl",
        init_method="tcp://localhost:29500",
        world_size=world_size,
        rank=rank,
        device_id=torch.device(f"cuda:{local_rank}"),
    )
    yield
    dist.destroy_process_group()


MultideviceTest

# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

# Run command:
# mpirun -np [num_devices] pytest tests/python/multidevice/[test_name].py --only-mpi -s

import os
import pytest
from typing import Iterable

import torch
import torch.distributed as dist

import nvfuser_direct as nvfuser

class MultideviceTest:
    def __init__(self):
        os.environ["NVIDIA_TF32_OVERRIDE"] = "0"

        self._communicator = nvfuser.multidevice.Communicator.instance()

        torch.cuda.set_device(self._communicator.local_rank())

        # This way, when individual tests create unsharded input, each rank
        # receives the same data.
        torch.manual_seed(0)

    @property
    def communicator(self):
        return self._communicator

    @property
    def size(self):
        return self._communicator.size()

    @property
    def rank(self):
        return self._communicator.rank()

    @property
    def local_size(self):
        return self._communicator.local_size()

    @property
    def local_rank(self):
        return self._communicator.local_rank()

    def shard_tensor_1d(self, t: torch.Tensor, dim: int, mesh) -> torch.Tensor:
        """Shard tensor along a single dimension (1D sharding only).

        Args:
            t: Tensor to shard (preferably on CPU for memory efficiency)
            dim: Dimension to shard along
            mesh: DeviceMesh to use for sharding

        Returns:
            Sharded tensor on current GPU device

        Example:
            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(num_devices))
            unsharded = torch.randn(num_devices, 4)
            sharded = self.shard_tensor_1d(unsharded, 0, mesh)
        """
        assert t.is_cpu, (
            "This is not strictly required but it's a general good practice "
            "for unit tests to create unsharded data on CPU to reduce GPU "
            "memory footprint."
        )
        return mesh.shard_tensor(t, dim).cuda(self.local_rank)

    def shard_tensor(self, t: torch.Tensor, tv) -> torch.Tensor:
        """Shard tensor using TensorView's parallelization and device mesh.

        Args:
            t: Tensor to shard (preferably on CPU for memory efficiency)
            tv: TensorView with device mesh and parallelization information

        Returns:
            Sharded tensor on current GPU device

        Example:
            with nvfuser.FusionDefinition() as fd:
                inp_tv = fd.define_tensor([-1, -1])
                inp_tv.set_device_mesh(mesh)
                inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)
                # ... rest of fusion

            unsharded = torch.randn(num_devices, 4)
            sharded = self.shard_tensor(unsharded, inp_tv)
        """
        assert t.is_cpu, (
            "This is not strictly required but it's a general good practice "
            "for unit tests to create unsharded data on CPU to reduce GPU "
            "memory footprint."
        )

        sharded = nvfuser.multidevice.shard_tensor(t, tv)
        return sharded.cuda(self.local_rank)

@pytest.fixture
def multidevice_test():
    fixture = MultideviceTest()
    yield fixture
    fixture.communicator.barrier()

def get_env(envs: Iterable[str], /, *, default: str) -> str:
    for env in envs:
        if value := os.environ.get(env):
            return value
    return default

# Set up the default process group for torch APIs like
# dist.device_mesh.init_device_mesh.
#
# This fixture is used by multi-GPU tests that use torch.distributed directly.


            with nvfuser.FusionDefinition() as fd:
                inp_tv = fd.define_tensor([-1, -1])
                inp_tv.set_device_mesh(mesh)
                inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)
                # ... rest of fusion

            unsharded = torch.randn(num_devices, 4)
            sharded = self.shard_tensor(unsharded, inp_tv)
        """
        assert t.is_cpu, (
            "This is not strictly required but it's a general good practice "
            "for unit tests to create unsharded data on CPU to reduce GPU "
            "memory footprint."
        )

        sharded = nvfuser.multidevice.shard_tensor(t, tv)
        return sharded.cuda(self.local_rank)

@pytest.fixture
def multidevice_test():
    fixture = MultideviceTest()
    yield fixture
    fixture.communicator.barrier()

def get_env(envs: Iterable[str], /, *, default: str) -> str:
    for env in envs:
        if value := os.environ.get(env):
            return value
    return default

# Set up the default process group for torch APIs like
# dist.device_mesh.init_device_mesh.
#
# This fixture is used by multi-GPU tests that use torch.distributed directly.
#
# I use "session" instead of "module" because
# https://github.com/pytorch/pytorch/issues/119196 reported race conditions
# when reinitializing process groups.
@pytest.fixture(scope="session")
def setup_default_process_group():
    # I avoided using nvfuser.Communicator to minimize fixture dependencies on


            unsharded = torch.randn(num_devices, 4)
            sharded = self.shard_tensor(unsharded, inp_tv)
        """
        assert t.is_cpu, (
            "This is not strictly required but it's a general good practice "
            "for unit tests to create unsharded data on CPU to reduce GPU "
            "memory footprint."
        )

        sharded = nvfuser.multidevice.shard_tensor(t, tv)
        return sharded.cuda(self.local_rank)

@pytest.fixture
def multidevice_test():
    fixture = MultideviceTest()
    yield fixture
    fixture.communicator.barrier()

def get_env(envs: Iterable[str], /, *, default: str) -> str:
    for env in envs:
        if value := os.environ.get(env):
            return value
    return default

# Set up the default process group for torch APIs like
# dist.device_mesh.init_device_mesh.
#
# This fixture is used by multi-GPU tests that use torch.distributed directly.
#
# I use "session" instead of "module" because
# https://github.com/pytorch/pytorch/issues/119196 reported race conditions
# when reinitializing process groups.
@pytest.fixture(scope="session")
def setup_default_process_group():
    # I avoided using nvfuser.Communicator to minimize fixture dependencies on
    # nvFuser. This makes the transition from legacy bindings to direct
    # bindings easier.
    rank = int(get_env(["OMPI_COMM_WORLD_RANK", "RANK"], default="0"))
    world_size = int(get_env(["OMPI_COMM_WORLD_SIZE", "WORLD_SIZE"], default="1"))
    local_rank = int(get_env(["OMPI_COMM_WORLD_LOCAL_RANK", "LOCAL_RANK"], default="0"))

    torch.cuda.set_device(local_rank)


    yield fixture
    fixture.communicator.barrier()

def get_env(envs: Iterable[str], /, *, default: str) -> str:
    for env in envs:
        if value := os.environ.get(env):
            return value
    return default

# Set up the default process group for torch APIs like
# dist.device_mesh.init_device_mesh.
#
# This fixture is used by multi-GPU tests that use torch.distributed directly.
#
# I use "session" instead of "module" because
# https://github.com/pytorch/pytorch/issues/119196 reported race conditions
# when reinitializing process groups.
@pytest.fixture(scope="session")
def setup_default_process_group():
    # I avoided using nvfuser.Communicator to minimize fixture dependencies on
    # nvFuser. This makes the transition from legacy bindings to direct
    # bindings easier.
    rank = int(get_env(["OMPI_COMM_WORLD_RANK", "RANK"], default="0"))
    world_size = int(get_env(["OMPI_COMM_WORLD_SIZE", "WORLD_SIZE"], default="1"))
    local_rank = int(get_env(["OMPI_COMM_WORLD_LOCAL_RANK", "LOCAL_RANK"], default="0"))

    torch.cuda.set_device(local_rank)

    # The default port as used by https://github.com/pytorch/pytorch/blob/45a8b5682eb69d865cbf68c7f2f689b56b4efd53/torch/csrc/distributed/c10d/TCPStore.hpp#L51.
    dist.init_process_group(
        backend="nccl",
        init_method="tcp://localhost:29500",
        world_size=world_size,
        rank=rank,
        device_id=torch.device(f"cuda:{local_rank}"),
    )
    yield
    dist.destroy_process_group()


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

# Run command:
# mpirun -np [num_devices] pytest tests/python/multidevice/[test_name].py --only-mpi -s

import os
import pytest
from typing import Iterable

import torch
import torch.distributed as dist

import nvfuser_direct as nvfuser

class MultideviceTest:
    def __init__(self):
        os.environ["NVIDIA_TF32_OVERRIDE"] = "0"

        self._communicator = nvfuser.multidevice.Communicator.instance()

        torch.cuda.set_device(self._communicator.local_rank())

        # This way, when individual tests create unsharded input, each rank
        # receives the same data.
        torch.manual_seed(0)

    @property
    def communicator(self):
        return self._communicator

    @property
    def size(self):
        return self._communicator.size()

    @property
    def rank(self):
        return self._communicator.rank()

    @property
    def local_size(self):
        return self._communicator.local_size()

    @property
    def local_rank(self):
        return self._communicator.local_rank()


communicator

import torch
import torch.distributed as dist

import nvfuser_direct as nvfuser

class MultideviceTest:
    def __init__(self):
        os.environ["NVIDIA_TF32_OVERRIDE"] = "0"

        self._communicator = nvfuser.multidevice.Communicator.instance()

        torch.cuda.set_device(self._communicator.local_rank())

        # This way, when individual tests create unsharded input, each rank
        # receives the same data.
        torch.manual_seed(0)

    @property
    def communicator(self):
        return self._communicator

    @property
    def size(self):
        return self._communicator.size()

    @property
    def rank(self):
        return self._communicator.rank()

    @property
    def local_size(self):
        return self._communicator.local_size()

    @property
    def local_rank(self):
        return self._communicator.local_rank()

    def shard_tensor_1d(self, t: torch.Tensor, dim: int, mesh) -> torch.Tensor:
        """Shard tensor along a single dimension (1D sharding only).


size

import nvfuser_direct as nvfuser

class MultideviceTest:
    def __init__(self):
        os.environ["NVIDIA_TF32_OVERRIDE"] = "0"

        self._communicator = nvfuser.multidevice.Communicator.instance()

        torch.cuda.set_device(self._communicator.local_rank())

        # This way, when individual tests create unsharded input, each rank
        # receives the same data.
        torch.manual_seed(0)

    @property
    def communicator(self):
        return self._communicator

    @property
    def size(self):
        return self._communicator.size()

    @property
    def rank(self):
        return self._communicator.rank()

    @property
    def local_size(self):
        return self._communicator.local_size()

    @property
    def local_rank(self):
        return self._communicator.local_rank()

    def shard_tensor_1d(self, t: torch.Tensor, dim: int, mesh) -> torch.Tensor:
        """Shard tensor along a single dimension (1D sharding only).

        Args:
            t: Tensor to shard (preferably on CPU for memory efficiency)
            dim: Dimension to shard along
            mesh: DeviceMesh to use for sharding


rank

    def __init__(self):
        os.environ["NVIDIA_TF32_OVERRIDE"] = "0"

        self._communicator = nvfuser.multidevice.Communicator.instance()

        torch.cuda.set_device(self._communicator.local_rank())

        # This way, when individual tests create unsharded input, each rank
        # receives the same data.
        torch.manual_seed(0)

    @property
    def communicator(self):
        return self._communicator

    @property
    def size(self):
        return self._communicator.size()

    @property
    def rank(self):
        return self._communicator.rank()

    @property
    def local_size(self):
        return self._communicator.local_size()

    @property
    def local_rank(self):
        return self._communicator.local_rank()

    def shard_tensor_1d(self, t: torch.Tensor, dim: int, mesh) -> torch.Tensor:
        """Shard tensor along a single dimension (1D sharding only).

        Args:
            t: Tensor to shard (preferably on CPU for memory efficiency)
            dim: Dimension to shard along
            mesh: DeviceMesh to use for sharding

        Returns:
            Sharded tensor on current GPU device


        torch.cuda.set_device(self._communicator.local_rank())

        # This way, when individual tests create unsharded input, each rank
        # receives the same data.
        torch.manual_seed(0)

    @property
    def communicator(self):
        return self._communicator

    @property
    def size(self):
        return self._communicator.size()

    @property
    def rank(self):
        return self._communicator.rank()

    @property
    def local_size(self):
        return self._communicator.local_size()

    @property
    def local_rank(self):
        return self._communicator.local_rank()

    def shard_tensor_1d(self, t: torch.Tensor, dim: int, mesh) -> torch.Tensor:
        """Shard tensor along a single dimension (1D sharding only).

        Args:
            t: Tensor to shard (preferably on CPU for memory efficiency)
            dim: Dimension to shard along
            mesh: DeviceMesh to use for sharding

        Returns:
            Sharded tensor on current GPU device

        Example:
            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(num_devices))
            unsharded = torch.randn(num_devices, 4)
            sharded = self.shard_tensor_1d(unsharded, 0, mesh)


        # receives the same data.
        torch.manual_seed(0)

    @property
    def communicator(self):
        return self._communicator

    @property
    def size(self):
        return self._communicator.size()

    @property
    def rank(self):
        return self._communicator.rank()

    @property
    def local_size(self):
        return self._communicator.local_size()

    @property
    def local_rank(self):
        return self._communicator.local_rank()

    def shard_tensor_1d(self, t: torch.Tensor, dim: int, mesh) -> torch.Tensor:
        """Shard tensor along a single dimension (1D sharding only).

        Args:
            t: Tensor to shard (preferably on CPU for memory efficiency)
            dim: Dimension to shard along
            mesh: DeviceMesh to use for sharding

        Returns:
            Sharded tensor on current GPU device

        Example:
            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(num_devices))
            unsharded = torch.randn(num_devices, 4)
            sharded = self.shard_tensor_1d(unsharded, 0, mesh)
        """
        assert t.is_cpu, (
            "This is not strictly required but it's a general good practice "
            "for unit tests to create unsharded data on CPU to reduce GPU "


    @property
    def communicator(self):
        return self._communicator

    @property
    def size(self):
        return self._communicator.size()

    @property
    def rank(self):
        return self._communicator.rank()

    @property
    def local_size(self):
        return self._communicator.local_size()

    @property
    def local_rank(self):
        return self._communicator.local_rank()

    def shard_tensor_1d(self, t: torch.Tensor, dim: int, mesh) -> torch.Tensor:
        """Shard tensor along a single dimension (1D sharding only).

        Args:
            t: Tensor to shard (preferably on CPU for memory efficiency)
            dim: Dimension to shard along
            mesh: DeviceMesh to use for sharding

        Returns:
            Sharded tensor on current GPU device

        Example:
            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(num_devices))
            unsharded = torch.randn(num_devices, 4)
            sharded = self.shard_tensor_1d(unsharded, 0, mesh)
        """
        assert t.is_cpu, (
            "This is not strictly required but it's a general good practice "
            "for unit tests to create unsharded data on CPU to reduce GPU "
            "memory footprint."
        )
        return mesh.shard_tensor(t, dim).cuda(self.local_rank)

    def shard_tensor(self, t: torch.Tensor, tv) -> torch.Tensor:
        """Shard tensor using TensorView's parallelization and device mesh.

        Args:
            t: Tensor to shard (preferably on CPU for memory efficiency)
            tv: TensorView with device mesh and parallelization information

        Returns:
            Sharded tensor on current GPU device

        Example:
            with nvfuser.FusionDefinition() as fd:
                inp_tv = fd.define_tensor([-1, -1])
                inp_tv.set_device_mesh(mesh)
                inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)
                # ... rest of fusion

            unsharded = torch.randn(num_devices, 4)
            sharded = self.shard_tensor(unsharded, inp_tv)


        Args:
            t: Tensor to shard (preferably on CPU for memory efficiency)
            dim: Dimension to shard along
            mesh: DeviceMesh to use for sharding

        Returns:
            Sharded tensor on current GPU device

        Example:
            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(num_devices))
            unsharded = torch.randn(num_devices, 4)
            sharded = self.shard_tensor_1d(unsharded, 0, mesh)
        """
        assert t.is_cpu, (
            "This is not strictly required but it's a general good practice "
            "for unit tests to create unsharded data on CPU to reduce GPU "
            "memory footprint."
        )
        return mesh.shard_tensor(t, dim).cuda(self.local_rank)

    def shard_tensor(self, t: torch.Tensor, tv) -> torch.Tensor:
        """Shard tensor using TensorView's parallelization and device mesh.

        Args:
            t: Tensor to shard (preferably on CPU for memory efficiency)
            tv: TensorView with device mesh and parallelization information

        Returns:
            Sharded tensor on current GPU device

        Example:
            with nvfuser.FusionDefinition() as fd:
                inp_tv = fd.define_tensor([-1, -1])
                inp_tv.set_device_mesh(mesh)
                inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)
                # ... rest of fusion

            unsharded = torch.randn(num_devices, 4)
            sharded = self.shard_tensor(unsharded, inp_tv)
        """
        assert t.is_cpu, (
            "This is not strictly required but it's a general good practice "
            "for unit tests to create unsharded data on CPU to reduce GPU "
            "memory footprint."
        )

        sharded = nvfuser.multidevice.shard_tensor(t, tv)
        return sharded.cuda(self.local_rank)

@pytest.fixture
def multidevice_test():
    fixture = MultideviceTest()
    yield fixture
    fixture.communicator.barrier()

def get_env(envs: Iterable[str], /, *, default: str) -> str:
    for env in envs:
        if value := os.environ.get(env):
            return value
    return default

# Set up the default process group for torch APIs like
# dist.device_mesh.init_device_mesh.
#
# This fixture is used by multi-GPU tests that use torch.distributed directly.


# SPDX-FileCopyrightText: Copyright (c) 2026-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest

import torch

import nvfuser_direct as nvfuser

# Unit tests for multidevice_test.shard_tensor.
class TestShardTensor:
    @pytest.mark.mpi
    def test_inner_split(self, multidevice_test):
        d = multidevice_test.size

        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            tv.set_device_mesh(mesh)
            tv.inner_split(0, d)
            tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        t_ref = torch.arange(3 * d, dtype=torch.float)
        t = multidevice_test.shard_tensor(t_ref, tv)
        torch.testing.assert_close(
            t.cpu(), t_ref.reshape(3, d)[:, multidevice_test.rank]
        )

    @pytest.mark.mpi
    def test_nonoutermost_split(self, multidevice_test):
        d = multidevice_test.size

        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            tv.set_device_mesh(mesh)
            tv.outer_split(0, 2)
            tv.outer_split(1, d)
            tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        t_ref = torch.arange(2 * d * 3, dtype=torch.float)
        t = multidevice_test.shard_tensor(t_ref, tv)
        torch.testing.assert_close(
            t.cpu(), t_ref.reshape(2, d, 3)[:, multidevice_test.rank, :].flatten()
        )

    @pytest.mark.mpi
    def test_parallel_reduction(self, multidevice_test):
        d = multidevice_test.size
        with nvfuser.FusionDefinition() as fd:
            inp_tv = fd.define_tensor([-1, -1])
            out_tv = fd.ops.sum(inp_tv, [1])
            fd.add_output(out_tv)

            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            for tv in [inp_tv, out_tv]:
                tv.set_device_mesh(mesh)
                tv.outer_split(1, d)
                tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        inp_ref = torch.arange(2 * d * 3, dtype=torch.float).reshape(2, d * 3)
        out_ref = inp_ref.sum([1])

        rank = multidevice_test.rank
        inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
        torch.testing.assert_close(inp.cpu(), inp_ref[:, rank * 3 : (rank + 1) * 3])
        out = multidevice_test.shard_tensor(out_ref, out_tv)
        torch.testing.assert_close(out.cpu(), out_ref)

    @pytest.mark.mpi
    def test_2d_sharded_matrix(self, multidevice_test):
        d = multidevice_test.size
        tp_size = 2
        if d % tp_size != 0:
            pytest.skip(
                f"Number of devices ({d}) must be divisible by tp_size ({tp_size})"
            )
        dp_size = d // tp_size

        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1, -1])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(
                torch.arange(d).reshape(dp_size, tp_size)
            )
            tv.set_device_mesh(mesh)
            tv.outer_split(0, dp_size)
            tv.axis(0).parallelize(nvfuser.ParallelType.mesh_y)
            tv.outer_split(2, tp_size)
            tv.axis(2).parallelize(nvfuser.ParallelType.mesh_x)

        rows_per_rank = 2
        cols_per_rank = 3
        t_ref = torch.arange(
            dp_size * rows_per_rank * tp_size * cols_per_rank, dtype=torch.float
        ).reshape(dp_size * rows_per_rank, tp_size * cols_per_rank)
        t = multidevice_test.shard_tensor(t_ref, tv)

        dp_rank = multidevice_test.rank // tp_size
        tp_rank = multidevice_test.rank % tp_size
        torch.testing.assert_close(
            t.cpu(),
            t_ref[
                dp_rank * rows_per_rank : (dp_rank + 1) * rows_per_rank,
                tp_rank * cols_per_rank : (tp_rank + 1) * cols_per_rank,
            ],
        )

    @pytest.mark.mpi
    def test_2d_sharded_vector(self, multidevice_test):
        d = multidevice_test.size
        tp_size = 2
        if d % tp_size != 0:
            pytest.skip(
                f"Number of devices ({d}) must be divisible by tp_size ({tp_size})"
            )
        dp_size = d // tp_size

        rows_per_rank = 2
        cols_per_rank = 3

        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(
                torch.arange(d).reshape(dp_size, tp_size)
            )
            tv.set_device_mesh(mesh)
            tv.inner_split(0, tp_size * cols_per_rank)
            tv.outer_split(0, dp_size)
            tv.axis(0).parallelize(nvfuser.ParallelType.mesh_y)
            tv.outer_split(2, tp_size)
            tv.axis(2).parallelize(nvfuser.ParallelType.mesh_x)

        t_ref = torch.arange(
            dp_size * rows_per_rank * tp_size * cols_per_rank, dtype=torch.float
        )
        t = multidevice_test.shard_tensor(t_ref, tv)

        dp_rank = multidevice_test.rank // tp_size
        tp_rank = multidevice_test.rank % tp_size
        torch.testing.assert_close(
            t.cpu(),
            t_ref.reshape(dp_size, rows_per_rank, tp_size, cols_per_rank)[
                dp_rank, :, tp_rank, :
            ].flatten(),
        )

    @pytest.mark.mpi
    def test_context_and_tensor_parallel(self, multidevice_test):
        d = multidevice_test.size
        tp_size = 2
        if d % tp_size != 0:
            pytest.skip(
                f"Number of devices ({d}) must be divisible by tp_size ({tp_size})"
            )
        cp_size = d // tp_size
        e = 5

        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1, -1, e])  # [b, s, e]
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(
                torch.arange(d).reshape(cp_size, tp_size)
            )
            tv.set_device_mesh(mesh)
            tv.outer_split(1, tp_size)
            tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)
            tv.outer_split(2, cp_size)
            tv.axis(2).parallelize(nvfuser.ParallelType.mesh_y)
            tv.reorder({2: 0, 1: 1})

        b, s = 2, tp_size * cp_size * 3
        t_ref = torch.arange(b * s * e, dtype=torch.float).reshape(b, s, e)
        t = multidevice_test.shard_tensor(t_ref, tv)

        cp_rank = multidevice_test.rank // tp_size
        tp_rank = multidevice_test.rank % tp_size
        torch.testing.assert_close(
            t.cpu(), t_ref.reshape(b, tp_size, cp_size, 3, e)[:, tp_rank, cp_rank, :, :]
        )

    @pytest.mark.mpi
    def test_expanded_broadcast(self, multidevice_test):
        d = multidevice_test.size
        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1, -1], contiguity=[True, None])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            tv.set_device_mesh(mesh)
            tv.outer_split(1, d)
            tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        # When d=2, t_ref =
        #   [[0, 0, 0, 0, 0, 0],
        #    [1, 1, 1, 1, 1, 1]]
        #
        # GPU 0 gets the left half, and GPU 1 gets the right half. So t =
        #   [[0, 0, 0],
        #    [1, 1, 1]]
        # regardless of which GPU.
        t_ref = torch.arange(2, dtype=torch.float).unsqueeze(-1).expand(-1, d * 3)
        t = multidevice_test.shard_tensor(t_ref, tv)
        torch.testing.assert_close(
            t.cpu(), torch.arange(2, dtype=torch.float).unsqueeze(-1).expand(-1, 3)
        )

    @pytest.mark.mpi
    def test_broadcast(self, multidevice_test):
        d = multidevice_test.size
        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([1, -1], contiguity=[None, True])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            tv.set_device_mesh(mesh)
            tv.outer_split(1, d)
            tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        # When d=2, t_ref = [[0, 1, 2, 3, 4, 5]]. GPU 0 gets [[0, 1, 2]], and GPU 1 gets [[3, 4, 5]].
        t_ref = torch.arange(d * 3, dtype=torch.float).unsqueeze(0)
        t = multidevice_test.shard_tensor(t_ref, tv)
        rank = multidevice_test.rank
        torch.testing.assert_close(t.cpu(), t_ref[:, rank * 3 : (rank + 1) * 3])


# SPDX-FileCopyrightText: Copyright (c) 2026-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest

import torch

import nvfuser_direct as nvfuser

# Unit tests for multidevice_test.shard_tensor.
class TestShardTensor:
    @pytest.mark.mpi
    def test_inner_split(self, multidevice_test):
        d = multidevice_test.size

        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            tv.set_device_mesh(mesh)
            tv.inner_split(0, d)
            tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        t_ref = torch.arange(3 * d, dtype=torch.float)
        t = multidevice_test.shard_tensor(t_ref, tv)
        torch.testing.assert_close(
            t.cpu(), t_ref.reshape(3, d)[:, multidevice_test.rank]
        )

    @pytest.mark.mpi
    def test_nonoutermost_split(self, multidevice_test):
        d = multidevice_test.size

        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            tv.set_device_mesh(mesh)
            tv.outer_split(0, 2)
            tv.outer_split(1, d)
            tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        t_ref = torch.arange(2 * d * 3, dtype=torch.float)
        t = multidevice_test.shard_tensor(t_ref, tv)
        torch.testing.assert_close(
            t.cpu(), t_ref.reshape(2, d, 3)[:, multidevice_test.rank, :].flatten()
        )


    @pytest.mark.mpi
    def test_inner_split(self, multidevice_test):
        d = multidevice_test.size

        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            tv.set_device_mesh(mesh)
            tv.inner_split(0, d)
            tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        t_ref = torch.arange(3 * d, dtype=torch.float)
        t = multidevice_test.shard_tensor(t_ref, tv)
        torch.testing.assert_close(
            t.cpu(), t_ref.reshape(3, d)[:, multidevice_test.rank]
        )

    @pytest.mark.mpi
    def test_nonoutermost_split(self, multidevice_test):
        d = multidevice_test.size

        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            tv.set_device_mesh(mesh)
            tv.outer_split(0, 2)
            tv.outer_split(1, d)
            tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        t_ref = torch.arange(2 * d * 3, dtype=torch.float)
        t = multidevice_test.shard_tensor(t_ref, tv)
        torch.testing.assert_close(
            t.cpu(), t_ref.reshape(2, d, 3)[:, multidevice_test.rank, :].flatten()
        )

    @pytest.mark.mpi
    def test_parallel_reduction(self, multidevice_test):
        d = multidevice_test.size
        with nvfuser.FusionDefinition() as fd:
            inp_tv = fd.define_tensor([-1, -1])
            out_tv = fd.ops.sum(inp_tv, [1])
            fd.add_output(out_tv)

            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            for tv in [inp_tv, out_tv]:
                tv.set_device_mesh(mesh)
                tv.outer_split(1, d)
                tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        inp_ref = torch.arange(2 * d * 3, dtype=torch.float).reshape(2, d * 3)
        out_ref = inp_ref.sum([1])

        rank = multidevice_test.rank
        inp = multidevice_test.shard_tensor(inp_ref, inp_tv)


    def test_nonoutermost_split(self, multidevice_test):
        d = multidevice_test.size

        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            tv.set_device_mesh(mesh)
            tv.outer_split(0, 2)
            tv.outer_split(1, d)
            tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        t_ref = torch.arange(2 * d * 3, dtype=torch.float)
        t = multidevice_test.shard_tensor(t_ref, tv)
        torch.testing.assert_close(
            t.cpu(), t_ref.reshape(2, d, 3)[:, multidevice_test.rank, :].flatten()
        )

    @pytest.mark.mpi
    def test_parallel_reduction(self, multidevice_test):
        d = multidevice_test.size
        with nvfuser.FusionDefinition() as fd:
            inp_tv = fd.define_tensor([-1, -1])
            out_tv = fd.ops.sum(inp_tv, [1])
            fd.add_output(out_tv)

            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            for tv in [inp_tv, out_tv]:
                tv.set_device_mesh(mesh)
                tv.outer_split(1, d)
                tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        inp_ref = torch.arange(2 * d * 3, dtype=torch.float).reshape(2, d * 3)
        out_ref = inp_ref.sum([1])

        rank = multidevice_test.rank
        inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
        torch.testing.assert_close(inp.cpu(), inp_ref[:, rank * 3 : (rank + 1) * 3])
        out = multidevice_test.shard_tensor(out_ref, out_tv)
        torch.testing.assert_close(out.cpu(), out_ref)

    @pytest.mark.mpi
    def test_2d_sharded_matrix(self, multidevice_test):
        d = multidevice_test.size
        tp_size = 2
        if d % tp_size != 0:
            pytest.skip(
                f"Number of devices ({d}) must be divisible by tp_size ({tp_size})"
            )
        dp_size = d // tp_size

        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1, -1])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(
                torch.arange(d).reshape(dp_size, tp_size)
            )
            tv.set_device_mesh(mesh)
            tv.outer_split(0, dp_size)


            inp_tv = fd.define_tensor([-1, -1])
            out_tv = fd.ops.sum(inp_tv, [1])
            fd.add_output(out_tv)

            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            for tv in [inp_tv, out_tv]:
                tv.set_device_mesh(mesh)
                tv.outer_split(1, d)
                tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        inp_ref = torch.arange(2 * d * 3, dtype=torch.float).reshape(2, d * 3)
        out_ref = inp_ref.sum([1])

        rank = multidevice_test.rank
        inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
        torch.testing.assert_close(inp.cpu(), inp_ref[:, rank * 3 : (rank + 1) * 3])
        out = multidevice_test.shard_tensor(out_ref, out_tv)
        torch.testing.assert_close(out.cpu(), out_ref)

    @pytest.mark.mpi
    def test_2d_sharded_matrix(self, multidevice_test):
        d = multidevice_test.size
        tp_size = 2
        if d % tp_size != 0:
            pytest.skip(
                f"Number of devices ({d}) must be divisible by tp_size ({tp_size})"
            )
        dp_size = d // tp_size

        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1, -1])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(
                torch.arange(d).reshape(dp_size, tp_size)
            )
            tv.set_device_mesh(mesh)
            tv.outer_split(0, dp_size)
            tv.axis(0).parallelize(nvfuser.ParallelType.mesh_y)
            tv.outer_split(2, tp_size)
            tv.axis(2).parallelize(nvfuser.ParallelType.mesh_x)

        rows_per_rank = 2
        cols_per_rank = 3
        t_ref = torch.arange(
            dp_size * rows_per_rank * tp_size * cols_per_rank, dtype=torch.float
        ).reshape(dp_size * rows_per_rank, tp_size * cols_per_rank)
        t = multidevice_test.shard_tensor(t_ref, tv)

        dp_rank = multidevice_test.rank // tp_size
        tp_rank = multidevice_test.rank % tp_size
        torch.testing.assert_close(
            t.cpu(),
            t_ref[
                dp_rank * rows_per_rank : (dp_rank + 1) * rows_per_rank,
                tp_rank * cols_per_rank : (tp_rank + 1) * cols_per_rank,
            ],
        )

    @pytest.mark.mpi
    def test_2d_sharded_vector(self, multidevice_test):
        d = multidevice_test.size
        tp_size = 2
        if d % tp_size != 0:
            pytest.skip(
                f"Number of devices ({d}) must be divisible by tp_size ({tp_size})"
            )
        dp_size = d // tp_size

        rows_per_rank = 2
        cols_per_rank = 3

        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(
                torch.arange(d).reshape(dp_size, tp_size)


            tv.axis(2).parallelize(nvfuser.ParallelType.mesh_x)

        rows_per_rank = 2
        cols_per_rank = 3
        t_ref = torch.arange(
            dp_size * rows_per_rank * tp_size * cols_per_rank, dtype=torch.float
        ).reshape(dp_size * rows_per_rank, tp_size * cols_per_rank)
        t = multidevice_test.shard_tensor(t_ref, tv)

        dp_rank = multidevice_test.rank // tp_size
        tp_rank = multidevice_test.rank % tp_size
        torch.testing.assert_close(
            t.cpu(),
            t_ref[
                dp_rank * rows_per_rank : (dp_rank + 1) * rows_per_rank,
                tp_rank * cols_per_rank : (tp_rank + 1) * cols_per_rank,
            ],
        )

    @pytest.mark.mpi
    def test_2d_sharded_vector(self, multidevice_test):
        d = multidevice_test.size
        tp_size = 2
        if d % tp_size != 0:
            pytest.skip(
                f"Number of devices ({d}) must be divisible by tp_size ({tp_size})"
            )
        dp_size = d // tp_size

        rows_per_rank = 2
        cols_per_rank = 3

        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(
                torch.arange(d).reshape(dp_size, tp_size)
            )
            tv.set_device_mesh(mesh)
            tv.inner_split(0, tp_size * cols_per_rank)
            tv.outer_split(0, dp_size)
            tv.axis(0).parallelize(nvfuser.ParallelType.mesh_y)
            tv.outer_split(2, tp_size)
            tv.axis(2).parallelize(nvfuser.ParallelType.mesh_x)

        t_ref = torch.arange(
            dp_size * rows_per_rank * tp_size * cols_per_rank, dtype=torch.float
        )
        t = multidevice_test.shard_tensor(t_ref, tv)

        dp_rank = multidevice_test.rank // tp_size
        tp_rank = multidevice_test.rank % tp_size
        torch.testing.assert_close(
            t.cpu(),
            t_ref.reshape(dp_size, rows_per_rank, tp_size, cols_per_rank)[
                dp_rank, :, tp_rank, :
            ].flatten(),
        )

    @pytest.mark.mpi
    def test_context_and_tensor_parallel(self, multidevice_test):
        d = multidevice_test.size
        tp_size = 2
        if d % tp_size != 0:
            pytest.skip(
                f"Number of devices ({d}) must be divisible by tp_size ({tp_size})"
            )
        cp_size = d // tp_size
        e = 5

        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1, -1, e])  # [b, s, e]
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(
                torch.arange(d).reshape(cp_size, tp_size)
            )
            tv.set_device_mesh(mesh)


            tv.outer_split(0, dp_size)
            tv.axis(0).parallelize(nvfuser.ParallelType.mesh_y)
            tv.outer_split(2, tp_size)
            tv.axis(2).parallelize(nvfuser.ParallelType.mesh_x)

        t_ref = torch.arange(
            dp_size * rows_per_rank * tp_size * cols_per_rank, dtype=torch.float
        )
        t = multidevice_test.shard_tensor(t_ref, tv)

        dp_rank = multidevice_test.rank // tp_size
        tp_rank = multidevice_test.rank % tp_size
        torch.testing.assert_close(
            t.cpu(),
            t_ref.reshape(dp_size, rows_per_rank, tp_size, cols_per_rank)[
                dp_rank, :, tp_rank, :
            ].flatten(),
        )

    @pytest.mark.mpi
    def test_context_and_tensor_parallel(self, multidevice_test):
        d = multidevice_test.size
        tp_size = 2
        if d % tp_size != 0:
            pytest.skip(
                f"Number of devices ({d}) must be divisible by tp_size ({tp_size})"
            )
        cp_size = d // tp_size
        e = 5

        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1, -1, e])  # [b, s, e]
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(
                torch.arange(d).reshape(cp_size, tp_size)
            )
            tv.set_device_mesh(mesh)
            tv.outer_split(1, tp_size)
            tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)
            tv.outer_split(2, cp_size)
            tv.axis(2).parallelize(nvfuser.ParallelType.mesh_y)
            tv.reorder({2: 0, 1: 1})

        b, s = 2, tp_size * cp_size * 3
        t_ref = torch.arange(b * s * e, dtype=torch.float).reshape(b, s, e)
        t = multidevice_test.shard_tensor(t_ref, tv)

        cp_rank = multidevice_test.rank // tp_size
        tp_rank = multidevice_test.rank % tp_size
        torch.testing.assert_close(
            t.cpu(), t_ref.reshape(b, tp_size, cp_size, 3, e)[:, tp_rank, cp_rank, :, :]
        )

    @pytest.mark.mpi
    def test_expanded_broadcast(self, multidevice_test):
        d = multidevice_test.size
        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1, -1], contiguity=[True, None])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            tv.set_device_mesh(mesh)
            tv.outer_split(1, d)
            tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        # When d=2, t_ref =
        #   [[0, 0, 0, 0, 0, 0],
        #    [1, 1, 1, 1, 1, 1]]
        #
        # GPU 0 gets the left half, and GPU 1 gets the right half. So t =
        #   [[0, 0, 0],
        #    [1, 1, 1]]


                torch.arange(d).reshape(cp_size, tp_size)
            )
            tv.set_device_mesh(mesh)
            tv.outer_split(1, tp_size)
            tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)
            tv.outer_split(2, cp_size)
            tv.axis(2).parallelize(nvfuser.ParallelType.mesh_y)
            tv.reorder({2: 0, 1: 1})

        b, s = 2, tp_size * cp_size * 3
        t_ref = torch.arange(b * s * e, dtype=torch.float).reshape(b, s, e)
        t = multidevice_test.shard_tensor(t_ref, tv)

        cp_rank = multidevice_test.rank // tp_size
        tp_rank = multidevice_test.rank % tp_size
        torch.testing.assert_close(
            t.cpu(), t_ref.reshape(b, tp_size, cp_size, 3, e)[:, tp_rank, cp_rank, :, :]
        )

    @pytest.mark.mpi
    def test_expanded_broadcast(self, multidevice_test):
        d = multidevice_test.size
        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1, -1], contiguity=[True, None])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            tv.set_device_mesh(mesh)
            tv.outer_split(1, d)
            tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        # When d=2, t_ref =
        #   [[0, 0, 0, 0, 0, 0],
        #    [1, 1, 1, 1, 1, 1]]
        #
        # GPU 0 gets the left half, and GPU 1 gets the right half. So t =
        #   [[0, 0, 0],
        #    [1, 1, 1]]
        # regardless of which GPU.
        t_ref = torch.arange(2, dtype=torch.float).unsqueeze(-1).expand(-1, d * 3)
        t = multidevice_test.shard_tensor(t_ref, tv)
        torch.testing.assert_close(
            t.cpu(), torch.arange(2, dtype=torch.float).unsqueeze(-1).expand(-1, 3)
        )

    @pytest.mark.mpi
    def test_broadcast(self, multidevice_test):
        d = multidevice_test.size
        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([1, -1], contiguity=[None, True])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            tv.set_device_mesh(mesh)
            tv.outer_split(1, d)
            tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        # When d=2, t_ref = [[0, 1, 2, 3, 4, 5]]. GPU 0 gets [[0, 1, 2]], and GPU 1 gets [[3, 4, 5]].
        t_ref = torch.arange(d * 3, dtype=torch.float).unsqueeze(0)
        t = multidevice_test.shard_tensor(t_ref, tv)
        rank = multidevice_test.rank
        torch.testing.assert_close(t.cpu(), t_ref[:, rank * 3 : (rank + 1) * 3])


            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            tv.set_device_mesh(mesh)
            tv.outer_split(1, d)
            tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        # When d=2, t_ref =
        #   [[0, 0, 0, 0, 0, 0],
        #    [1, 1, 1, 1, 1, 1]]
        #
        # GPU 0 gets the left half, and GPU 1 gets the right half. So t =
        #   [[0, 0, 0],
        #    [1, 1, 1]]
        # regardless of which GPU.
        t_ref = torch.arange(2, dtype=torch.float).unsqueeze(-1).expand(-1, d * 3)
        t = multidevice_test.shard_tensor(t_ref, tv)
        torch.testing.assert_close(
            t.cpu(), torch.arange(2, dtype=torch.float).unsqueeze(-1).expand(-1, 3)
        )

    @pytest.mark.mpi
    def test_broadcast(self, multidevice_test):
        d = multidevice_test.size
        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([1, -1], contiguity=[None, True])
            fd.add_output(tv)

            mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
            tv.set_device_mesh(mesh)
            tv.outer_split(1, d)
            tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

        # When d=2, t_ref = [[0, 1, 2, 3, 4, 5]]. GPU 0 gets [[0, 1, 2]], and GPU 1 gets [[3, 4, 5]].
        t_ref = torch.arange(d * 3, dtype=torch.float).unsqueeze(0)
        t = multidevice_test.shard_tensor(t_ref, tv)
        rank = multidevice_test.rank
        torch.testing.assert_close(t.cpu(), t_ref[:, rank * 3 : (rank + 1) * 3])