Diff Coverage

        Returns:
            tuple: (local_args, local_kwargs, cache_values)
        """
        args, kwargs = _normalize_chunk_view_args(*args, **kwargs)
        input_tensor, chunks, dim = args
        input_shape = input_tensor.shape

        local_args = (input_tensor.to_local(), chunks, dim)
        local_kwargs = {}

        cache_values = [input_tensor.layout, chunks, dim, input_shape]
        return local_args, local_kwargs, cache_values

    def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
        """
        Infer output layouts for ChunkView operator.

            )

        mapping = alias_map[dim]
        if isinstance(mapping, (list, tuple)):
            is_sharded = any(m != "None" for m in mapping)
        else:
            is_sharded = mapping != "None"

        if is_sharded:

platform = get_platform()


def _normalize_histc_args(x, bins=100, min=0, max=0):  # pylint: disable=W0622
    return (x, bins, min, max), {}


class HistcExtDistributedOp(DistributedOp):
    """

        Returns:
            tuple: (local_args, local_kwargs, cache_values)
        """
        args, kwargs = _normalize_histc_args(*args, **kwargs)
        x, bins, min_val, max_val = args
        local_args = (x.to_local(), bins, min_val, max_val)
        local_kwargs = {}
        cache_values = [x.layout, bins, min_val, max_val]
        return local_args, local_kwargs, cache_values

    def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
        """
        Infer output layout for HistcExt operator.

                f"For {self.op_name}, bins should be a positive integer, "
                f"but got {bins}"
            )
        if not isinstance(min_val, (int, float)):
            raise ValueError(
                f"For {self.op_name}, min should be a number, "
                f"but got {type(min_val).__name__}"
            )
        if not isinstance(max_val, (int, float)):
            raise ValueError(
                f"For {self.op_name}, max should be a number, "
                f"but got {type(max_val).__name__}"
            )
        if min_val > max_val:

def _normalize_inplace_scatter_value_args(input_tensor, dim, index, value):
    """Normalize InplaceScatterValue arguments to positional args + empty kwargs."""
    return (input_tensor, dim, index, value), {}


class InplaceScatterValueDistributedOp(DistributedOp):
    """Distributed implementation for InplaceScatterValue operator."""

        Returns:
            tuple: (local_args, local_kwargs, cache_values)
        """
        args, kwargs = _normalize_inplace_scatter_value_args(*args, **kwargs)
        input_tensor, dim, index, value = args

        local_args = (
            input_tensor.to_local() if hasattr(input_tensor, '_layout') else input_tensor,
            dim,
            index.to_local() if hasattr(index, '_layout') else index,
            value,

            index.to_local() if hasattr(index, '_layout') else index,
            value,
        )

        cache_values = [
            input_tensor.layout if hasattr(input_tensor, '_layout') else None,
            index.layout if hasattr(index, '_layout') else None,
            dim,
        ]
        local_kwargs = kwargs
        return local_args, local_kwargs, cache_values

    def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
        """
        Infer output layout for InplaceScatterValue operator.

                f"For {self.op_name}, dim {original_dim} is out of bounds for tensor with {ndim} dims"
            )

        if len(input_map) != len(index_map):
            raise ValueError(
                f"For {self.op_name}, input and index must have the same number of dimensions, "
                f"but got input rank={len(input_map)}, index rank={len(index_map)}"
            )

            out_layout.set_partial_by_dev_axis(axis_alias, op)


def _normalize_matmul_ext_args(x, w):
    return (x, w), {}


class MatMulExtDistributedOp(DistributedOp):
    """Distributed implementation for MatMul operator."""

        Returns:
            tuple: (local_args, local_kwargs, cache_values) where local_args contains
                local tensors for x and w; cache_values contains [x_layout, w_layout].
        """
        args, kwargs = _normalize_matmul_ext_args(*args, **kwargs)
        x_tensor, w_tensor = args[0], args[1]
        local_args = (x_tensor.to_local(), w_tensor.to_local())
        local_kwargs = {}
        cache_values = [x_tensor.layout, w_tensor.layout]
        return local_args, local_kwargs, cache_values

    def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
        """
        Infer output layout for MatMul operator (output = x @ w).

        """
        x_layout = cache_values[0]
        w_layout = cache_values[1]
        if not x_layout or not w_layout:
            raise ValueError(
                f"For {self.op_name}, x_layout: {x_layout}, w_layout: {w_layout}"
            )
        x_mesh_shape = x_layout.mesh_shape
        w_mesh_shape = w_layout.mesh_shape

            )
        x_mesh_shape = x_layout.mesh_shape
        w_mesh_shape = w_layout.mesh_shape
        if x_mesh_shape != w_mesh_shape:
            raise ValueError(
                f"For {self.op_name}, inputs must have same mesh_shape, "
                f"but got x: {x_mesh_shape} and w: {w_mesh_shape}"
            )

        w_map = w_layout.alias_tensor_map
        contract_dim = len(x_map) - 1
        w_contract_dim = len(w_map) - 2
        if x_map[contract_dim] != w_map[w_contract_dim]:
            raise ValueError(
                f"For {self.op_name}, contracting dimensions must have same layout, "
                f"but got x: {x_map[contract_dim]} and w: {w_map[w_contract_dim]}"
            )

        return ((out_layout,), None)


def _normalize_matmul_args(x, w, transpose_a=False, transpose_b=False):
    return (x, w, transpose_a, transpose_b), {}


class MatMulDistributedOp(DistributedOp):
    """Distributed implementation for MatMul operator."""

        Returns:
            tuple: (local_args, local_kwargs, cache_values) where local_args contains
                local tensors for x and w; cache_values contains [x_layout, w_layout, transpose_a, transpose_b].
        """
        args, kwargs = _normalize_matmul_args(*args, **kwargs)
        x_tensor, w_tensor, transpose_a, transpose_b = args
        local_args = (x_tensor.to_local(), w_tensor.to_local())
        local_kwargs = {}
        cache_values = [x_tensor.layout, w_tensor.layout, transpose_a, transpose_b]
        return local_args, local_kwargs, cache_values

    def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
        """
        Infer output layout for MatMul operator (output = x @ w, with possible transpose).

        transpose_a = cache_values[2]
        transpose_b = cache_values[3]

        if not x_layout or not w_layout:
            raise ValueError(
                f"For {self.op_name}, x_layout: {x_layout}, w_layout: {w_layout}"
            )

        x_mesh_shape = x_layout.mesh_shape

        x_mesh_shape = x_layout.mesh_shape
        w_mesh_shape = w_layout.mesh_shape
        if x_mesh_shape != w_mesh_shape:
            raise ValueError(
                f"For {self.op_name}, inputs must have same mesh_shape, "
                f"but got x: {x_mesh_shape} and w: {w_mesh_shape}"
            )

            w_contract_dim = len(w_map) - 2  # Second to last dimension

        # Validate contracting dimensions
        if x_map[x_contract_dim] != w_map[w_contract_dim]:
            raise ValueError(
                f"For {self.op_name}, contracting dimensions must have same layout, "
                f"but got x: {x_map[x_contract_dim]} and w: {w_map[w_contract_dim]}"
            )

        # Set partial status
        if x_map[x_contract_dim] != "None":
            if isinstance(x_map[x_contract_dim], tuple):
                for axis in x_map[x_contract_dim]:
                    out_layout.set_partial_by_dev_axis(axis, 'sum')
            else:
                out_layout.set_partial_by_dev_axis(x_map[x_contract_dim], 'sum')

        return ((out_layout,), None)

        return output_layout


def _normalize_batch_matmul_ext_args(x, w):
    return (x, w), {}


class BatchMatMulExtDistributedOp(BaseBatchMatMulDistributedOp):
    """Distributed implementation for BatchMatMulExt operator."""

        Returns:
            tuple: (local_args, local_kwargs, cache_values) where local_args contains
                local tensors for x and w; cache_values contains [x_layout, w_layout].
        """
        args, kwargs = _normalize_batch_matmul_ext_args(*args, **kwargs)
        x_tensor, w_tensor = args[0], args[1]
        local_args = (x_tensor.to_local(), w_tensor.to_local())
        local_kwargs = {}
        cache_values = [x_tensor.layout, w_tensor.layout]
        return local_args, local_kwargs, cache_values

    def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
        """
        Infer output layout for BatchMatMulExt operator (output = x @ w).

        x_layout = cache_values[0]
        w_layout = cache_values[1]

        if not x_layout or not w_layout:
            raise ValueError(
                f"For {self.op_name}, x_layout: {x_layout}, w_layout: {w_layout}"
            )

        if x_layout.mesh_shape != w_layout.mesh_shape:
            raise ValueError(
                f"For {self.op_name}, inputs must have same mesh_shape, "
                f"but got x: {x_layout.mesh_shape} and w: {w_layout.mesh_shape}"
            )

        # contracting dims
        x_contract = x_map[-1]
        w_contract = w_map[-2]
        if x_contract != w_contract:
            raise ValueError(
                f"For {self.op_name}, contracting (M) dim layouts must match, "
                f"but got x: {x_contract} and w: {w_contract}"
            )

        return ((self._build_output_layout(x_layout, w_layout, merged_batch, x_n, w_p, x_contract),), None)


def _normalize_batch_matmul_args(x, w, transpose_a=False, transpose_b=False):
    return (x, w, transpose_a, transpose_b), {}


class BatchMatMulDistributedOp(BaseBatchMatMulDistributedOp):
    """Distributed implementation for BatchMatMul operator."""

            tuple: (local_args, local_kwargs, cache_values) where local_args contains
                local tensors for x and w; cache_values contains
                [x_layout, w_layout, transpose_a, transpose_b].
        """
        args, kwargs = _normalize_batch_matmul_args(*args, **kwargs)
        x_tensor, w_tensor, transpose_a, transpose_b = args
        local_args = (x_tensor.to_local(), w_tensor.to_local())
        local_kwargs = {}
        cache_values = [x_tensor.layout, w_tensor.layout, transpose_a, transpose_b]
        return local_args, local_kwargs, cache_values

    def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
        """
        Infer output layout for BatchMatMul operator (output = x @ w, with possible transpose).

        transpose_a = cache_values[2]
        transpose_b = cache_values[3]

        if not x_layout or not w_layout:
            raise ValueError(
                f"For {self.op_name}, x_layout: {x_layout}, w_layout: {w_layout}"
            )

        if x_layout.mesh_shape != w_layout.mesh_shape:
            raise ValueError(
                f"For {self.op_name}, inputs must have same mesh_shape, "
                f"but got x: {x_layout.mesh_shape} and w: {w_layout.mesh_shape}"
            )

            w_contract = w_map[-2]
            w_p = w_map[-1]

        if x_contract != w_contract:
            raise ValueError(
                f"For {self.op_name}, contracting (M) dim layouts must match, "
                f"but got x: {x_contract} and w: {w_contract}"
            )

        Returns:
            tuple: (local_args, local_kwargs, cache_values)
        """
        args, kwargs = _normalize_scatter_update_args(*args, **kwargs)
        x, indices, updates = args

        local_args = (
            x.to_local() if hasattr(x, '_layout') else x,
            indices.to_local() if hasattr(indices, '_layout') else indices,
            updates.to_local() if hasattr(updates, '_layout') else updates,
        )

            indices.to_local() if hasattr(indices, '_layout') else indices,
            updates.to_local() if hasattr(updates, '_layout') else updates,
        )

        cache_values = [
            x.layout if hasattr(x, '_layout') else None,
            indices.layout if hasattr(indices, '_layout') else None,
            updates.layout if hasattr(updates, '_layout') else None,
        ]
        local_kwargs = kwargs
        return local_args, local_kwargs, cache_values

    def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
        """
        Infer output layout for ScatterUpdate operator.

        """
        input_layout, indices_layout, updates_layout = cache_values

        if input_layout is None:
            raise ValueError(
                f"For {self.op_name}, input layout should not be None"
            )

        if indices_layout is None:

        # Partial inputs are intentionally allowed. The output inherits Partial status
        # from the input layout (see lines below), making this a Partial-preserving op.
        if not self._allow_partial_inputs:
            self._check_partial_inputs([input_layout])

        self._validate_strategy(input_layout, indices_layout, updates_layout)

        output_layout = Layout(

        indices_map = indices_layout.alias_tensor_map
        updates_map = updates_layout.alias_tensor_map

        if not input_map:
            raise ValueError(
                f"For {self.op_name}, input tensor map should not be empty"
            )

        if input_map[0] != "None":

        valid_layouts = [lyt for lyt in layouts if lyt is not None]

        if not valid_layouts:
            raise ValueError(
                f"For {self.op_name}, stack requires at least one input DTensor."
            )

        # Reference layout to validate consistency across all input tensors