Diff Coverage

        self._check_ce_op_without_loss_parallel_context(op_name, args)

        if not is_loss_parallel_op(op_name):
            raise RuntimeError(
                f"Operator {op_name} does not contain parallel layout infer func. "
                f"DTensor dispatch requires explicit layout inference registration. "
                f"Please register a distributed operator for '{op_name}' or use local tensors."
            )

            logits = gathered_args[0]
            targets = gathered_args[1]
            if isinstance(logits, Tensor) and isinstance(targets, Tensor):
                if logits.ndim > 2 and targets.ndim > 1 and targets.ndim == logits.ndim - 1:
                    vocab_size = logits.shape[-1]
                    gathered_args[0] = logits.reshape(-1, vocab_size)
                    gathered_args[1] = targets.reshape(-1)

        return op_call(*gathered_args, **gathered_kwargs)

    def _dispatch_layout_infer(

        # Use alias_tensor_map to support StridedShard multi-axis mappings.
        alias_map = input_layout.alias_tensor_map
        mapping = alias_map[axis]
        if isinstance(mapping, tuple):
            is_sharded = any(m != "None" for m in mapping)
        else:
            is_sharded = mapping != "None"

        if is_sharded:

from .parallel_ops import DistributedOp


def _normalize_atleast_1d_args(*tensors):
    return tensors, {}


class Atleast1DDistributedOp(DistributedOp):
    """Distributed implementation for torch.atleast_1d."""

        Returns:
            tuple: (local_args, local_kwargs, cache_values)
        """
        args, kwargs = _normalize_atleast_1d_args(*args, **kwargs)
        tensors = args

        local_args = tuple(
            t.to_local() if hasattr(t, 'to_local') else t
            for t in tensors
        )
        local_kwargs = {}

        cache_values = [
            t.layout if hasattr(t, 'layout') else None
            for t in tensors
        ]

        return local_args, local_kwargs, cache_values

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

        # If there are multiple inputs, return a tuple of Layouts.
        if len(output_layouts) == 1:
            return ((output_layouts[0],), None)

        return (tuple(output_layouts), None)

        if output_layouts[0] is None:
            return py_output
        # Use from_local_with_layout to reuse the pre-computed Layout
        # instead of rebuilding it in DTensor.__init_data__.
        return DTensor.from_local_with_layout(
            py_output, output_layouts[0]
        )

        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:

from .parallel_ops import DistributedOp


def _normalize_expand_args(input_tensor, *sizes):
    return (input_tensor, *sizes), {}


def _normalize_expand_as_args(input_tensor, target_tensor):
    return (input_tensor, target_tensor), {}


class ExpandDistributedOp(DistributedOp):
    """Distributed implementation for torch.Tensor.expand."""

        Returns:
            tuple: (local_args, local_kwargs, cache_values)
        """
        args, kwargs = _normalize_expand_args(*args, **kwargs)
        input_tensor = args[0]
        sizes = tuple(args[1:])
        local_args = (input_tensor.to_local(), *sizes)
        cache_values = [input_tensor.layout, input_tensor.shape, sizes]
        return local_args, {}, cache_values

    @staticmethod
    def _validate_input_layouts(
        cache_values: list,

            )

        input_shape = cache_values[1] if len(cache_values) > 1 else None
        if not isinstance(input_shape, tuple):
            raise ValueError(
                f"For {op_name}, input_shape should be a tuple, "
                f"but got {type(input_shape)}."
            )

            )

        sizes = cache_values[2] if len(cache_values) > 2 else None
        if sizes is None or len(sizes) < 1:
            raise ValueError(
                f"For {op_name}, sizes should be a non-empty tuple of ints, "
                f"but got {sizes}."
            )
        for i, sz in enumerate(sizes):

                f"but got {sizes}."
            )
        for i, sz in enumerate(sizes):
            if not isinstance(sz, int):
                raise ValueError(
                    f"For {op_name}, elements in sizes should be int, "
                    f"but got {type(sz)} at position {i}."
                )

                        f"got mapping {in_alias_map[i]}."
                    )
                output_map.append("None")
            else:
                raise ValueError(
                    f"For {self.op_name}, cannot expand dimension {i} "
                    f"from size {input_size} to {requested_size}."
                )

        Returns:
            tuple: (local_args, local_kwargs, cache_values)
        """
        args, kwargs = _normalize_expand_as_args(*args, **kwargs)
        input_tensor = args[0]
        target_tensor = args[1]
        local_args = (input_tensor.to_local(), target_tensor.to_local())
        cache_values = [input_tensor.layout, input_tensor.shape, target_tensor.shape]
        return local_args, {}, cache_values

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

platform = get_platform()


def _normalize_histc_args(x, bins=100, min_val=0, max_val=0):
    return (x, bins, min_val, max_val), {}


class HistcExtDistributedOp(DistributedOp):
    """

            tuple: (local_args, local_kwargs, cache_values)
        """
        # Map external API parameter names (min, max) to internal names to avoid
        # shadowing Python builtins.
        if "min" in kwargs:
            kwargs["min_val"] = kwargs.pop("min")
        if "max" in kwargs:
            kwargs["max_val"] = kwargs.pop("max")
        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)}"
            )

from .parallel_ops import DistributedOp


def _normalize_masked_scatter_args(input_tensor, mask, source):
    return (input_tensor, mask, source), {}


class MaskedScatterDistributedOp(DistributedOp):
    """Distributed implementation for torch.Tensor.masked_scatter."""

        Returns:
            tuple: (local_args, local_kwargs, cache_values)
        """
        args, kwargs = _normalize_masked_scatter_args(*args, **kwargs)
        input_tensor, mask, source = args[0], args[1], args[2]
        local_args = (
            input_tensor.to_local(),
            mask.to_local(),
            source.to_local(),
        )
        local_kwargs = {}
        cache_values = [
            input_tensor.layout,
            mask.layout,
            source.layout,
        ]
        return local_args, local_kwargs, cache_values

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

            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(), transpose_a, transpose_b)
        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(), transpose_a, transpose_b)
        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}"
            )

def _normalize_mhc_pre_clamp_sinkhorn_args(*args, **kwargs):
    """Normalize npu_mhc_pre_clamp_sinkhorn arguments."""
    values = dict(_MHC_PRE_CLAMP_DEFAULTS)
    if len(args) > len(_MHC_PRE_CLAMP_ARG_NAMES):
        raise TypeError(
            f"npu_mhc_pre_clamp_sinkhorn expected at most {len(_MHC_PRE_CLAMP_ARG_NAMES)} arguments"
        )
    for name, value in zip(_MHC_PRE_CLAMP_ARG_NAMES, args):
        values[name] = value
    for name, value in kwargs.items():
        if name not in _MHC_PRE_CLAMP_ARG_NAMES:
            raise TypeError(f"npu_mhc_pre_clamp_sinkhorn got an unexpected keyword argument '{name}'")
        if name in _MHC_PRE_CLAMP_ARG_NAMES[:len(args)]:
            raise TypeError(f"npu_mhc_pre_clamp_sinkhorn got multiple values for argument '{name}'")
        values[name] = value
    missing = [name for name in _MHC_PRE_CLAMP_ARG_NAMES[:4] if name not in values]
    if missing:
        raise TypeError(f"npu_mhc_pre_clamp_sinkhorn missing required arguments: {missing}")
    return tuple(values[name] for name in _MHC_PRE_CLAMP_ARG_NAMES), {}


# Validation rules table for npu_mhc_pre_sinkhorn
# Key: tensor_map length (format identifier)

        #
        # kv_seq_split_num > 1 is blocked by a guard in
        # _compute_adjusted_sparse_params before reaching this function,
        # so local_kv_len == global_kv_len is guaranteed here.
        local_q_len = query.shape[seq_dim_idx]
        local_kv_len = key.shape[seq_dim_idx]

        if sparse_mode in (SPARSE_DEFAULT_MASK, SPARSE_BAND):
            new_pre_tokens = pre_tokens
            new_next_tokens = next_tokens

        if sparse_mode in (SPARSE_DEFAULT_MASK, SPARSE_BAND):
            new_pre_tokens = pre_tokens
            new_next_tokens = next_tokens
        else:
            new_pre_tokens = local_kv_len
            new_next_tokens = 0

        new_sparse_mode = SPARSE_BAND if sparse_mode != SPARSE_DEFAULT_MASK else sparse_mode
        update_mode = SPARSE_MODE_UPDATE_MAP[sparse_mode]

        new_sparse_mode = SPARSE_BAND if sparse_mode != SPARSE_DEFAULT_MASK else sparse_mode
        update_mode = SPARSE_MODE_UPDATE_MAP[sparse_mode]

        if update_mode == LEFT_UP_TO_LEFT_UP:
            offset = -split_id * local_q_len
            new_pre_tokens = new_pre_tokens + offset
            new_next_tokens = new_next_tokens - offset
        elif update_mode == LEFT_UP_TO_RIGHT_DOWN:
            offset = local_kv_len - (split_id + 1) * local_q_len
            new_pre_tokens = new_pre_tokens + offset
            new_next_tokens = new_next_tokens - offset
        elif update_mode == RIGHT_DOWN_TO_RIGHT_DOWN:
            offset = (split_num - split_id - 1) * local_q_len
            new_pre_tokens = new_pre_tokens + offset
            new_next_tokens = new_next_tokens - offset

        return new_sparse_mode, new_pre_tokens, new_next_tokens

        is_dynamic: bool,
    ) -> Tuple[int, int, int]:
        """Compute adjusted sparse parameters based on dynamic or static shape."""
        if is_dynamic:
            if kv_seq_split_num > 1:
                raise NotImplementedError(
                    f"For {self.op_name}, dynamic shape with KV sequence sharding "
                    f"(kv_seq_split_num={kv_seq_split_num}) is not yet supported. "
                    f"The dynamic path currently uses local KV length directly, "
                    f"while the static path multiplies by kv_seq_split_num to obtain "

            # Scalar parameters from pyboost may arrive as Tensor objects.
            # Arithmetic and comparison operations on such Tensor scalars
            # would trigger device kernel calls, so convert them to Python
            # native types early.
            head_num = int(self._to_python_scalar(p_head_num))
            keep_prob = self._to_python_scalar(p_keep_prob)
            scale_value = self._to_python_scalar(p_scale_value)
            pre_tokens = int(self._to_python_scalar(p_pre_tokens))
            next_tokens = int(self._to_python_scalar(p_next_tokens))
            inner_precise = int(self._to_python_scalar(p_inner_precise))
            sparse_mode = int(self._to_python_scalar(p_sparse_mode))

            # Ensure the runtime input_layout matches the cached value used
            # for sharding derivation and validation.
            runtime_input_layout = _resolve_input_layout(
                self._to_python_scalar(p_input_layout)
            )
            if runtime_input_layout != input_layout:
                raise ValueError(
                    f"For {self.op_name}, runtime input_layout {runtime_input_layout!r} "
                    f"does not match the cached input_layout {input_layout!r} "
                    f"used for sharding inference. This may indicate an incorrect "
                    f"dispatcher cache key."

                )

            is_varlen = input_layout == "TND" and actual_seq_qlen is not None
            self._validate_attn_mask(attn_mask, sparse_mode, input_layout, is_varlen)
            FlashAttentionScoreDistributedOp._validate_real_shift_configuration(
                real_shift, sparse_mode)

            split_info = self._get_split_info(query_layout, input_layout)
            head_split_num = split_info["head"]

                return FlashAttentionScoreDistributedOp._truncate_result(result)

            adjusted_head_num = self._adjust_head_num(head_num, head_split_num)

            (adjusted_sparse_mode, adjusted_pre_tokens, adjusted_next_tokens,
             adjusted_actual_seq_qlen, adjusted_actual_seq_kvlen) = self._apply_seq_split_adjustments(
                query, key, query_layout, key_layout, input_layout,
                sparse_mode, pre_tokens, next_tokens,
                actual_seq_qlen, actual_seq_kvlen,

                actual_seq_qlen, actual_seq_kvlen,
                seq_split_num, lb_split_id, lb_split_num,
            )

            result = func(
                query, key, value,
                real_shift, drop_mask, padding_mask, attn_mask, prefix,
                adjusted_actual_seq_qlen, adjusted_actual_seq_kvlen,
                int(adjusted_head_num), keep_prob, scale_value,

                int(adjusted_pre_tokens), int(adjusted_next_tokens), inner_precise,
                p_input_layout, int(adjusted_sparse_mode),
            )

            return FlashAttentionScoreDistributedOp._truncate_result(result)

        return _expanded_impl

    def _apply_seq_split_adjustments(  # pylint: disable=too-many-arguments,too-many-locals

        Returns:
            Tuple of (adjusted_sparse_mode, adjusted_pre_tokens, adjusted_next_tokens,
                      adjusted_actual_seq_qlen, adjusted_actual_seq_kvlen).
        """
        adjusted_sparse_mode = sparse_mode
        adjusted_pre_tokens = pre_tokens
        adjusted_next_tokens = next_tokens
        adjusted_actual_seq_qlen = actual_seq_qlen
        adjusted_actual_seq_kvlen = actual_seq_kvlen

        if seq_split_num > 1 or lb_split_id is not None:
            dynamic_info = self._get_dynamic_shape_info(query, key, input_layout)
            is_dynamic = dynamic_info.get('is_dynamic', False)

            if lb_split_id is not None:
                if lb_split_num is None:
                    raise ValueError("lb_split_num must not be None when lb_split_id is set")
                split_id = lb_split_id
                seq_split_num = lb_split_num
            else:
                split_id = self._get_split_id(query_layout, input_layout)
            seq_dim_idx = self._get_seq_dim_idx(self._layout_dims.get(input_layout, {}))

            if seq_dim_idx is None:
                raise ValueError(
                    f"Cannot infer seq/total dim for input_layout={input_layout}"
                )

            kv_seq_split_num = 1
            if key_layout is not None:
                kv_split_info = self._get_split_info(key_layout, input_layout)
                kv_seq_split_num = kv_split_info["seq"]

            self._check_seq_sharding_compatibility(
                query_layout, key_layout, input_layout,
                seq_dim_idx, seq_split_num, kv_seq_split_num
            )

            (adjusted_sparse_mode,
             adjusted_pre_tokens,
             adjusted_next_tokens) = self._compute_adjusted_sparse_params(
                query, key,
                sparse_mode, pre_tokens, next_tokens,

                split_id, seq_split_num, seq_dim_idx,
                kv_seq_split_num, is_dynamic,
            )

            if input_layout == "TND":
                (adjusted_sparse_mode,
                 adjusted_pre_tokens,
                 adjusted_next_tokens,
                 adjusted_actual_seq_qlen,
                 adjusted_actual_seq_kvlen) = self._adjust_tnd_layout_params(

                        kv_seq_split_num=kv_seq_split_num, is_dynamic=is_dynamic,
                    ),
                )

        return (adjusted_sparse_mode, adjusted_pre_tokens, adjusted_next_tokens,
                adjusted_actual_seq_qlen, adjusted_actual_seq_kvlen)

    def _get_seq_dim_idx(self, dims: dict) -> Optional[int]:
        """Get the sequence dimension index."""

    """Normalize RmsNorm args to positional form.

    MindSpore Primitive RmsNorm receives (x, gamma, epsilon) as positional arguments.
    """
    return (x, gamma, epsilon), {}


def _normalize_layernorm_args(input_tensor, normalized_shape, weight=None, bias=None, eps=1e-5):
    """Normalize layer_norm args to positional form.

    torch.nn.functional.layer_norm(input_tensor, normalized_shape, weight=None, bias=None, eps=1e-5)
    has no keyword-only parameters, so everything stays positional.
    """
    return (input_tensor, normalized_shape, weight, bias, eps), {}


class NormDistributedOp(DistributedOp):
    """Distributed implementation for RmsNorm operator."""

        Returns:
            tuple: (local_args, local_kwargs, cache_values)
        """
        args, kwargs = _normalize_rmsnorm_args(*args, **kwargs)
        x, gamma, epsilon = args
        local_args = (x.to_local(), gamma.to_local(), epsilon)
        local_kwargs = {}
        cache_values = [x.layout, gamma.layout]
        return local_args, local_kwargs, cache_values

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

            raise ValueError(f"{self.op_name} inputs must have same mesh_shape")
        x_alias_map = x_layout.alias_tensor_map
        gamma_alias_map = gamma_layout.alias_tensor_map
        if len(gamma_alias_map) > len(x_alias_map):
            raise ValueError(
                f"For {self.op_name}, gamma ndim {len(gamma_alias_map)} cannot exceed "
                f"input ndim {len(x_alias_map)}."
            )
        begin_norm_axis = len(x_alias_map) - len(gamma_alias_map)

        Returns:
            tuple: (local_args, local_kwargs, cache_values)
        """
        args, kwargs = _normalize_layernorm_args(*args, **kwargs)
        input_tensor, normalized_shape, weight, bias, eps = args

        # Normalize normalized_shape: int → (int,), list → tuple
        if isinstance(normalized_shape, int):
            normalized_shape = (normalized_shape,)
        elif isinstance(normalized_shape, list):
            normalized_shape = tuple(normalized_shape)

        local_args = [
            input_tensor.to_local(),
            normalized_shape,
            weight.to_local() if weight is not None and hasattr(weight, 'to_local') else weight,
            bias.to_local() if bias is not None and hasattr(bias, 'to_local') else bias,

            weight.to_local() if weight is not None and hasattr(weight, 'to_local') else weight,
            bias.to_local() if bias is not None and hasattr(bias, 'to_local') else bias,
            eps,
        ]
        local_kwargs = {}

        cache_values = [input_tensor.layout, normalized_shape]
        return tuple(local_args), local_kwargs, cache_values

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

        if not self._allow_partial_inputs:
            self._check_partial_inputs([input_layout])

        if normalized_shape is None:
            raise ValueError(f"{self.op_name} requires normalized_shape.")

        if not isinstance(normalized_shape, tuple):
            raise ValueError(f"normalized_shape must be int, list, or tuple, got {type(normalized_shape)}")

    Returns:
        tuple: (positional_args_tuple, empty_kwargs_dict)
    """
    return (
        query, key, value, head_num, input_layout,
        pse, padding_mask, atten_mask,
        scale, keep_prob, pre_tockens, next_tockens,
        inner_precise, prefix, actual_seq_qlen, actual_seq_kvlen,

        Returns:
            tuple: (local_args, local_kwargs, cache_values)
        """
        args, kwargs = _normalize_npu_fusion_attention_args(*args, **kwargs)
        query, key, value = args[0], args[1], args[2]
        input_layout = args[4]

        local_args = (
            query.to_local() if hasattr(query, '_layout') else query,
            key.to_local() if hasattr(key, '_layout') else key,
            value.to_local() if hasattr(value, '_layout') else value,
        ) + args[3:]
        local_kwargs = {}

        query_layout = query.layout if hasattr(query, "_layout") else None
        key_layout = key.layout if hasattr(key, "_layout") else None
        value_layout = value.layout if hasattr(value, "_layout") else None

        cache_values = [query_layout, key_layout, value_layout, input_layout]
        return local_args, local_kwargs, cache_values

    def _is_dynamic_shape(self, tensor: Tensor, dim: int) -> bool:
        """Check if tensor has dynamic shape at given dimension."""
        try:

        #
        # kv_seq_split_num > 1 is blocked by a guard in
        # _compute_adjusted_sparse_params before reaching this function,
        # so local_kv_len == global_kv_len is guaranteed here.
        local_q_len = query.shape[seq_dim_idx]
        local_kv_len = key.shape[seq_dim_idx]

        if sparse_mode in (SPARSE_DEFAULT_MASK, SPARSE_BAND):
            new_pre_tockens = pre_tockens
            new_next_tockens = next_tockens

        if sparse_mode in (SPARSE_DEFAULT_MASK, SPARSE_BAND):
            new_pre_tockens = pre_tockens
            new_next_tockens = next_tockens
        else:
            new_pre_tockens = local_kv_len
            new_next_tockens = 0

        new_sparse_mode = SPARSE_BAND if sparse_mode != SPARSE_DEFAULT_MASK else sparse_mode
        update_mode = SPARSE_MODE_UPDATE_MAP[sparse_mode]

        new_sparse_mode = SPARSE_BAND if sparse_mode != SPARSE_DEFAULT_MASK else sparse_mode
        update_mode = SPARSE_MODE_UPDATE_MAP[sparse_mode]

        if update_mode == LEFT_UP_TO_LEFT_UP:
            offset = -split_id * local_q_len
            new_pre_tockens = new_pre_tockens + offset
            new_next_tockens = new_next_tockens - offset
        elif update_mode == LEFT_UP_TO_RIGHT_DOWN:
            offset = local_kv_len - (split_id + 1) * local_q_len
            new_pre_tockens = new_pre_tockens + offset
            new_next_tockens = new_next_tockens - offset
        elif update_mode == RIGHT_DOWN_TO_RIGHT_DOWN:
            offset = (split_num - split_id - 1) * local_q_len
            new_pre_tockens = new_pre_tockens + offset
            new_next_tockens = new_next_tockens - offset

        return new_sparse_mode, new_pre_tockens, new_next_tockens

        )

        attention_out_layout = copy.deepcopy(query_layout)
        if attention_out_layout.placements is None and attention_out_layout.tensor_map is not None:
            attention_out_layout.tensor_map_to_placement()

        softmax_layout = self._infer_softmax_layout_by_input_layout(
            query_layout, input_layout_str, ""
        )

        is_dynamic: bool,
    ) -> Tuple[int, int, int]:
        """Compute adjusted sparse parameters based on dynamic or static shape."""
        if is_dynamic:
            if kv_seq_split_num > 1:
                raise NotImplementedError(
                    f"For {self.op_name}, dynamic shape with KV sequence sharding "
                    f"(kv_seq_split_num={kv_seq_split_num}) is not yet supported. "
                    f"The dynamic path currently uses local KV length directly, "
                    f"while the static path multiplies by kv_seq_split_num to obtain "

        query_layout = cache_values[0]
        key_layout = cache_values[1]

        if query_layout is None:
            return None

        def _expanded_impl(  # pylint: disable=R0913
            query,
            key,

            dynamic_info = self._get_dynamic_shape_info(query, key, input_layout)
            is_dynamic = dynamic_info.get('is_dynamic', False)

            if lb_split_id is not None:
                if lb_split_num is None:
                    raise ValueError(
                        "lb_split_num must not be None when lb_split_id is set"
                    )
                split_id = lb_split_id
                seq_split_num = lb_split_num

platform = get_platform()


def _normalize_one_hot_ext_args(indices, num_classes, on_value, off_value, axis):
    return (indices, num_classes, on_value, off_value, axis), {}


class OneHotExtDistributedOp(DistributedOp):
    """Distributed implementation for OneHotExt operator."""

        Returns:
            tuple: (local_args, local_kwargs, cache_values)
        """
        args, kwargs = _normalize_one_hot_ext_args(*args, **kwargs)
        indices, num_classes, on_value, off_value, axis = args

        indices_local = indices.to_local()
        on_value_local = on_value.to_local() if hasattr(on_value, '_layout') else on_value
        off_value_local = off_value.to_local() if hasattr(off_value, '_layout') else off_value

        on_value_layout = on_value.layout if hasattr(on_value, '_layout') else None
        off_value_layout = off_value.layout if hasattr(off_value, '_layout') else None

        local_args = (indices_local, num_classes, on_value_local, off_value_local, axis)
        local_kwargs = {}
        cache_values = [indices.layout, on_value_layout, off_value_layout, num_classes, axis]
        return local_args, local_kwargs, cache_values

    # pylint: disable=W0237
    def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
        """

        axis = self._validate_axis(axis)

        in_tensor_map = indices_layout.tensor_map
        if not in_tensor_map:
            raise ValueError(
                f"For {self.op_name}, indices tensor_map is empty."
            )

        self._validate_multi_dim_restriction(in_tensor_map, axis, indices_layout)

def _get_alias_shard_set(dim_alias):
    if isinstance(dim_alias, str):
        return {dim_alias} if dim_alias != "None" else set()
    return set(dim_alias)


class OuterDistributedOp(DistributedOp):
    """Distributed implementation for torch.outer."""

        dim = cache_values[1]
        keepdim = cache_values[2]

        if x_layout is None or x_layout.mesh_shape is None:
            raise ValueError(
                f"For {self.op_name}, input layout cannot be None."
            )

        # Check partial inputs

            )

        # Check partial inputs
        if not self._allow_partial_inputs:
            self._check_partial_inputs([x_layout])

        if dim is not None and not isinstance(dim, (int, tuple, list)):
            raise TypeError(
                f"For {self.op_name}, the `dim` argument should be `None`, `int`, "

            return self._handle_all_axis_reduce(x_layout, keepdim)

        # Case 2: dim is an empty tuple/list — reduce no dimensions, output layout equals input.
        if isinstance(dim, (tuple, list)) and len(dim) == 0:
            output_layout = Layout(
                mesh_shape=x_layout.mesh_shape,
                alias_name=x_layout.alias_name,
                rank_list=x_layout.rank_list
            )
            return output_layout(*x_layout.alias_tensor_map)

        # Case 3: dim is int, tuple, or list with at least one element.
        output_layout = Layout(
            mesh_shape=x_layout.mesh_shape,

            )
        # torch.max(x) global reduction: only pass the tensor so the call
        # remains torch.max(local_x), not torch.max(local_x, None, False).
        if dim is None:
            local_args = (input_tensor.to_local(),)
        else:
            local_args = (input_tensor.to_local(), dim, keepdim)
        local_kwargs = {}