diff --git a/hyper_parallel/core/shard/_op_dispatch.py b/hyper_parallel/core/shard/_op_dispatch.py
index 78953cda6e9f5a6bdd724dcb06bf05fe11365e3d..a364d27ede66941a63a03ea370f691b9f02643db 100644
--- a/hyper_parallel/core/shard/_op_dispatch.py
+++ b/hyper_parallel/core/shard/_op_dispatch.py
@@ -238,7 +238,9 @@ class OpDispatcher:
         else:
             all_args = (input_layouts, extra_args)
             output_layout = distribute_op.infer_layout(*all_args)
-            op_impl = distribute_op.get_expand_impl(func, output_layout, input_layouts, extra_args)
+            op_impl = getattr(
+                distribute_op, "get_expand_impl", lambda *args, **kwargs: None
+                )(func, output_layout, input_layouts, extra_args)
             op_layout_cache[cache_key] = (output_layout, op_impl)
 
         if op_impl is None:
diff --git a/hyper_parallel/core/shard/ops/parallel_one_hot_ext.py b/hyper_parallel/core/shard/ops/parallel_one_hot_ext.py
new file mode 100644
index 0000000000000000000000000000000000000000..dd3e1b97925e257114a2446955b73393cee792a0
--- /dev/null
+++ b/hyper_parallel/core/shard/ops/parallel_one_hot_ext.py
@@ -0,0 +1,422 @@
+# Copyright 2026 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""
+Distributed implementation for OneHotExt operator.
+"""
+
+import mindspore as ms
+from mindspore import ops, Tensor
+
+from hyper_parallel.core.layout import Layout
+from hyper_parallel.platform import get_platform
+from .parallel_ops import DistributedOp
+
+
+platform = get_platform()
+
+
+class OneHotExtDistributedOp(DistributedOp):
+    """Distributed implementation for OneHotExt operator."""
+
+    def infer_layout(self, layouts, extra_args):
+        """
+        Infer output layout for OneHotExt.
+
+        Principle:
+            OneHot operator generates a one-hot encoded tensor by inserting a new dimension
+            at the specified axis position. The new dimension size equals num_classes and
+            is always replicated (not sharded).
+
+        Sharding Strategy:
+            - Supported sharding on indices tensor:
+              * For 1D indices: any dimension can be sharded with axis in [-1, 0, 1]
+              * For multi-dimensional indices (rank > 1):
+                - axis must be -1
+                - Only the first dimension can be sharded (data parallel only)
+                - Other dimensions must be fully replicated
+            - Non-indices inputs (num_classes, on_value, off_value, axis) must be fully replicated
+            - Indices tensor must NOT be in partial state (requires complete index values)
+
+        Layout Transformation:
+            - Input indices layout: (d0, d1, ..., dn)
+            - Output layout: (d0, d1, ..., d_axis, -1, d_axis+1, ..., dn)
+              where -1 represents the replicated one-hot dimension
+
+        Args:
+            layouts (tuple): Tuple containing input layouts (indices, num_classes, on_value, off_value, axis).
+            extra_args (tuple): Additional arguments containing [num_classes, on_value, off_value, axis].
+
+        Returns:
+            Layout: Output layout with one-hot dimension inserted at specified axis.
+
+        Raises:
+            ValueError: If indices layout is None, indices is in partial state,
+                       or sharding strategy violates constraints.
+        """
+        if not layouts:
+            return None
+
+        indices_layout = layouts[0]
+        if indices_layout is None or indices_layout.mesh_shape is None:
+            raise ValueError(f"{self.op_name}: indices layout cannot be None")
+
+        if indices_layout.is_partial():
+            raise ValueError(
+                f"{self.op_name}: indices cannot be in partial state. "
+                f"Indices must contain complete index values for OneHot operation."
+            )
+
+        num_classes = self._get_num_classes(extra_args)
+        self._validate_num_classes(num_classes)
+
+        axis = self._get_axis(extra_args)
+
+        in_tensor_map = indices_layout.tensor_map
+        if not in_tensor_map:
+            raise ValueError(f"{self.op_name}: indices tensor_map is empty")
+
+        self._validate_multi_dim_restriction(in_tensor_map, axis, indices_layout)
+        self._validate_inputs_layouts(layouts)
+
+        out_tensor_map = self._infer_output_tensor_map(in_tensor_map, axis)
+        out_layout = self._create_layout_from_tensor_map(indices_layout, out_tensor_map)
+        return out_layout
+
+    def get_expand_impl(self, func, output_layout, layouts, extra_args):
+        """
+        Get expanded implementation for OneHotExt operator.
+
+        Requirements:
+            - Indices must be int64; otherwise raise TypeError.
+            - num_classes must be int and >= -1.
+
+        When num_classes=-1 and indices is sharded:
+            Perform allreduce(max) across sharded mesh axes to infer a global depth value,
+            then call the original operator with a Python int depth.
+
+        Notes:
+            - MindSpore communication allreduce does NOT support int64 input for this backend.
+              We cast the scalar max value to int32 for communication and convert back to Python int.
+            - Some backends do not reliably handle 0-D tensors for allreduce; we unsqueeze to 1-D
+              before communication and squeeze back afterwards.
+
+        Args:
+            original_op: Original operator function.
+            output_layout: Inferred output layout (unused here).
+            input_layouts: List of input layouts [indices_layout, ...].
+            extra_args: Additional arguments [num_classes, on_value, off_value, axis].
+
+        Returns:
+            Callable: Expanded operator implementation or None if no expansion needed.
+        """
+        del output_layout  # kept for signature compatibility; not used
+
+        indices_layout = layouts[0] if layouts else None
+        if indices_layout is None:
+            return None
+
+        sharded_axes = self._get_sharded_axes(indices_layout)
+        if not sharded_axes:
+            return None
+
+        original_op = func
+        rank = platform.get_rank()
+        reduce_max = ops.ReduceMax(keep_dims=False)
+
+        def expanded_one_hot(indices, num_classes, on_value, off_value, axis):
+            self._validate_num_classes(num_classes)
+            self._validate_indices_dtype(indices)
+
+            if num_classes != -1:
+                return original_op(indices, num_classes, on_value, off_value, axis)
+
+            local_max = reduce_max(indices, ())
+            if not isinstance(local_max, Tensor):
+                local_max = Tensor(local_max, ms.int64)
+
+            # Ensure safe cast int64 -> int32 for communication.
+            # Depth is derived from max(indices) + 1, so values outside int32 are invalid here.
+            local_max_host = int(local_max.asnumpy())
+            if local_max_host > 2147483647:
+                raise ValueError(
+                    f"{self.op_name}: indices max value {local_max_host} exceeds int32 range, "
+                    f"cannot infer num_classes via allreduce(max)."
+                )
+
+            zero_dim = local_max.ndim == 0
+            local_max_i32 = ops.cast(local_max, ms.int32)
+
+            if zero_dim:
+                local_max_i32 = ops.expand_dims(local_max_i32, 0)
+
+            global_max_i32 = local_max_i32
+            for axis_name in sharded_axes:
+                group = indices_layout.get_comm_group_by_axis(axis_name, rank)
+                global_max_i32 = platform.differentiable_all_reduce(
+                    global_max_i32, "max", group
+                )
+
+            if zero_dim:
+                global_max_i32 = ops.squeeze(global_max_i32, 0)
+
+            depth = int(global_max_i32.asnumpy()) + 1
+            return original_op(indices, depth, on_value, off_value, axis)
+
+        return expanded_one_hot
+
+    def _get_num_classes(self, extra_args):
+        """
+        Extract num_classes from extra arguments.
+
+        Args:
+            extra_args (list/tuple): Additional arguments [num_classes, on_value, off_value, axis].
+
+        Returns:
+            int: num_classes value, default -1 when not provided.
+        """
+        if isinstance(extra_args, (list, tuple)) and len(extra_args) >= 1:
+            num_classes = extra_args[0]
+            if isinstance(num_classes, int):
+                return num_classes
+        return -1
+
+    def _validate_num_classes(self, num_classes):
+        """
+        Validate num_classes parameter.
+
+        Raises:
+            TypeError: If num_classes is not int.
+            ValueError: If num_classes < -1.
+        """
+        if not isinstance(num_classes, int):
+            raise TypeError(
+                f"{self.op_name}: num_classes must be int, but got {type(num_classes).__name__}"
+            )
+        if num_classes < -1:
+            raise ValueError(
+                f"{self.op_name}: num_classes must be >= -1, but got {num_classes}"
+            )
+
+    def _validate_indices_dtype(self, indices):
+        """
+        Validate indices dtype. Backend requires int64.
+
+        Raises:
+            TypeError: If indices dtype is not int64.
+        """
+        if indices.dtype != ms.int64:
+            raise TypeError(
+                f"{self.op_name}: indices dtype must be int64, but got {indices.dtype}"
+            )
+
+    def _get_sharded_axes(self, layout):
+        """
+        Get all device axes that are used for sharding.
+
+        Args:
+            layout: Tensor layout.
+
+        Returns:
+            list: List of device axis names that are sharded.
+        """
+        sharded_axes = set()
+
+        if layout is None or layout.alias_tensor_map is None:
+            return []
+
+        for dim_alias in layout.alias_tensor_map:
+            if dim_alias == "None":
+                continue
+
+            if isinstance(dim_alias, tuple):
+                for axis_name in dim_alias:
+                    if axis_name != "None":
+                        sharded_axes.add(axis_name)
+            else:
+                sharded_axes.add(dim_alias)
+
+        return list(sharded_axes)
+
+    def _get_axis(self, extra_args):
+        """
+        Extract axis parameter from extra arguments.
+
+        Args:
+            extra_args (tuple): Additional arguments [num_classes, on_value, off_value, axis].
+
+        Returns:
+            int: Axis value for one-hot dimension insertion, default is -1.
+        """
+        if isinstance(extra_args, (list, tuple)) and len(extra_args) >= 4:
+            axis = extra_args[3]
+            if isinstance(axis, int):
+                return self._validate_axis(axis)
+        return -1
+
+    def _validate_axis(self, axis):
+        """
+        Validate axis parameter.
+
+        Args:
+            axis (int): Axis value to validate.
+
+        Returns:
+            int: Validated axis value.
+
+        Raises:
+            TypeError: If axis is not int.
+            ValueError: If axis is out of range[-1, 1].
+        """
+        if not isinstance(axis, int):
+            raise TypeError(
+                f"{self.op_name}: axis must be int, but got {type(axis).__name__}"
+            )
+
+        if axis > 1 or axis < -1:
+            raise ValueError(f"{self.op_name}: axis {axis} is out of range[-1, 1]")
+
+        return axis
+
+    def _validate_multi_dim_restriction(self, in_tensor_map, axis, indices_layout):
+        """
+        Validate restriction for multi-dimensional inputs.
+
+        Args:
+            in_tensor_map (tuple): Input tensor map.
+            axis (int): Axis value.
+            indices_layout (Layout): Indices layout.
+
+        Raises:
+            ValueError: If input dimension > 1 and axis is not -1, or strategy is not data parallel.
+        """
+        in_rank = len(in_tensor_map)
+        if in_rank <= 1:
+            return
+
+        if axis != -1:
+            raise ValueError(
+                f"{self.op_name}: when input dimension is > 1, axis must be -1, but got {axis}"
+            )
+
+        alias_map = indices_layout.alias_tensor_map
+        for i in range(1, len(alias_map)):
+            if alias_map[i] != "None":
+                raise ValueError(
+                    f"{self.op_name}: when input dimension is > 1, strategy must be data parallel, "
+                    f"but dimension {i} is sharded on '{alias_map[i]}'"
+                )
+
+    def _validate_inputs_layouts(self, layouts):
+        """
+        Validate that non-indices inputs are fully replicated.
+
+        Args:
+            layouts (tuple): Input layouts to validate.
+
+        Raises:
+            ValueError: If any non-indices input is sharded.
+        """
+        for layout in layouts[1:]:
+            if layout is None:
+                continue
+            alias_map = layout.alias_tensor_map
+            if alias_map and any(x != "None" for x in alias_map):
+                raise ValueError(
+                    f"{self.op_name}: non-indices inputs must be replicated, but got {alias_map}"
+                )
+
+    def _infer_output_tensor_map(self, in_tensor_map, axis):
+        """
+        Infer output tensor map by inserting one-hot dimension at specified axis.
+
+        Args:
+            in_tensor_map (tuple): Input tensor map from indices layout.
+            axis (int): Axis position for one-hot dimension insertion.
+
+        Returns:
+            tuple: Output tensor map with -1 inserted at axis position.
+
+        Raises:
+            ValueError: If axis is out of valid range.
+        """
+        in_rank = len(in_tensor_map)
+
+        if axis in (-1, in_rank):
+            insert_pos = in_rank
+        else:
+            insert_pos = axis
+
+        if insert_pos < 0 or insert_pos > in_rank:
+            raise ValueError(
+                f"{self.op_name}: axis {axis} is out of range for input with rank {in_rank}"
+            )
+
+        out_tensor_map = list(in_tensor_map)
+        out_tensor_map.insert(insert_pos, -1)
+        return tuple(out_tensor_map)
+
+    def _create_layout_from_tensor_map(self, base_layout, out_tensor_map):
+        """
+        Create output layout from tensor map.
+
+        Args:
+            base_layout (Layout): Base layout to derive mesh configuration.
+            out_tensor_map (tuple): Output tensor map.
+
+        Returns:
+            Layout: Complete output layout with tensor map and alias tensor map set.
+        """
+        out_layout = Layout(
+            mesh_shape=base_layout.mesh_shape,
+            alias_name=base_layout.alias_name,
+            rank_list=base_layout.rank_list,
+        )
+
+        out_layout.set_tensor_map(out_tensor_map)
+        out_layout.set_alias_tensor_map(
+            self._tensor_map_to_alias_tensor_map(base_layout, out_tensor_map)
+        )
+        out_layout.update_compact_str()
+        return out_layout
+
+    def _tensor_map_to_alias_tensor_map(self, base_layout, tensor_map):
+        """
+        Convert numeric tensor map to alias tensor map.
+
+        Args:
+            base_layout (Layout): Base layout providing alias names.
+            tensor_map (tuple): Numeric tensor map to convert.
+
+        Returns:
+            tuple: Alias tensor map with device axis names.
+        """
+        alias_name = base_layout.alias_name
+        alias_tensor_map = []
+
+        for dim in tensor_map:
+            if dim == -1:
+                alias_tensor_map.append("None")
+                continue
+
+            if isinstance(dim, tuple):
+                names = tuple(
+                    alias_name[len(alias_name) - 1 - d] for d in dim if d != -1
+                )
+                alias_tensor_map.append(names if names else "None")
+                continue
+
+            alias_tensor_map.append(alias_name[len(alias_name) - 1 - dim])
+
+        return tuple(alias_tensor_map)
diff --git a/hyper_parallel/core/shard/ops/yaml/one_hot_ext_ops.yaml b/hyper_parallel/core/shard/ops/yaml/one_hot_ext_ops.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..c62552e5934ae3f3f100965f5491ee8e3b8da3a6
--- /dev/null
+++ b/hyper_parallel/core/shard/ops/yaml/one_hot_ext_ops.yaml
@@ -0,0 +1,4 @@
+OneHotExt:
+  dist_op_name: _one_hot_ext_dist_op
+  distributed_op_class: OneHotExtDistributedOp
+  distributed_op_file: parallel_one_hot_ext
\ No newline at end of file
diff --git a/tests/mindspore/st/shard/one_hot_ext_shard_in_python.py b/tests/mindspore/st/shard/one_hot_ext_shard_in_python.py
new file mode 100644
index 0000000000000000000000000000000000000000..b0501ec9184adaedbb630f20b25673991234d3c4
--- /dev/null
+++ b/tests/mindspore/st/shard/one_hot_ext_shard_in_python.py
@@ -0,0 +1,319 @@
+# Copyright 2026 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""test one hot ext shard in python"""
+
+import numpy as np
+
+import mindspore as ms
+import mindspore.communication.management as D
+from mindspore import nn, Tensor
+from mindspore import mint
+from hyper_parallel import Layout, shard, DTensor
+from tests.mindspore.st.shard.utils import global_to_local, local_to_global
+
+
+def setup_module():
+    ms.set_device("Ascend")
+    D.init()
+
+
+class OneHotExtNet(nn.Cell):
+    """OneHot network using mint.nn.functional.one_hot"""
+
+    def __init__(self, relu_strategy=None):
+        super().__init__()
+        self.relu = nn.ReLU()
+        if relu_strategy is not None:
+            shard(self.relu, {"forward": {"input": relu_strategy}})
+
+    def construct(self, tensor, num_classes):
+        out = mint.nn.functional.one_hot(tensor, num_classes)
+        out = self.relu(out)
+        out = out + 1
+        return out
+
+
+base_mesh_shape = (2, 2, 2)
+base_alias_name = ("dp", "cp", "mp")
+base_rank_list = list(range(8))
+
+
+def _build_layout():
+    return Layout(base_mesh_shape, base_alias_name, base_rank_list)
+
+
+def _make_global_from_dp_shards(first_shard, second_shard):
+    """
+    Build global indices tensor by concatenating two dp shards along dim 0.
+
+    Note:
+        For layout("dp") on mesh_shape (2,2,2), dp has 2 shards.
+        Global tensor = concat([dp0_shard, dp1_shard], axis=0).
+    """
+    return Tensor(np.concatenate([first_shard, second_shard], axis=0), ms.int64)
+
+
+def _get_dp_id(rank):
+    """
+    For mesh_shape (2,2,2) with rank_list 0..7:
+        ranks 0..3 -> dp=0
+        ranks 4..7 -> dp=1
+    """
+    return rank // 4
+
+
+def test_one_hot_ext_data_parallel_1d_int64_1():
+    """
+    Feature: mint.nn.functional.one_hot in python shard.
+    Description: Test with 1D tensor, data parallel on tensor dim 0 (indices int64).
+    Expectation: Run success.
+    """
+    ms.set_seed(1)
+    np.random.seed(1)
+    a, num_classes = 64, 32
+
+    tensor = Tensor(np.random.randint(0, num_classes, size=(a,)).astype(np.int64), ms.int64)
+
+    standalone_net = OneHotExtNet()
+    standalone_output = standalone_net(tensor, num_classes)
+
+    layout = _build_layout()
+    tensor_layout = layout("dp")
+    onehot_out_layout = layout("dp", "None")
+
+    tensor_local = global_to_local(tensor, tensor_layout)
+
+    parallel_net = OneHotExtNet(relu_strategy=(onehot_out_layout,))
+    parallel_output = parallel_net(tensor_local, num_classes)
+
+    parallel_output = local_to_global(parallel_output)
+    assert np.allclose(
+        standalone_output.asnumpy(), parallel_output.asnumpy(), 1e-3, 1e-3
+    )
+
+
+def test_one_hot_ext_data_parallel_2d_int64_2():
+    """
+    Feature: mint.nn.functional.one_hot in python shard.
+    Description: Test with 2D tensor, data parallel on tensor dim 0 (indices int64).
+    Expectation: Run success.
+    """
+    ms.set_seed(1)
+    np.random.seed(1)
+    a, b, num_classes = 16, 32, 64
+
+    tensor = Tensor(np.random.randint(0, num_classes, size=(a, b)).astype(np.int64), ms.int64)
+
+    standalone_net = OneHotExtNet()
+    standalone_output = standalone_net(tensor, num_classes)
+
+    layout = _build_layout()
+    tensor_layout = layout("dp", "None")
+    onehot_out_layout = layout("dp", "None", "None")
+
+    tensor_local = global_to_local(tensor, tensor_layout)
+
+    parallel_net = OneHotExtNet(relu_strategy=(onehot_out_layout,))
+    parallel_output = parallel_net(tensor_local, num_classes)
+
+    parallel_output = local_to_global(parallel_output)
+    assert np.allclose(
+        standalone_output.asnumpy(), parallel_output.asnumpy(), 1e-3, 1e-3
+    )
+
+
+def test_one_hot_ext_replicate_all_int64_3():
+    """
+    Feature: mint.nn.functional.one_hot in python shard.
+    Description: Test with full replication (indices int64).
+    Expectation: Run success.
+    """
+    ms.set_seed(1)
+    np.random.seed(1)
+    a, num_classes = 64, 32
+
+    tensor = Tensor(np.random.randint(0, num_classes, size=(a,)).astype(np.int64), ms.int64)
+
+    standalone_net = OneHotExtNet()
+    standalone_output = standalone_net(tensor, num_classes)
+
+    layout = _build_layout()
+    tensor_layout = layout("None")
+    onehot_out_layout = layout("None", "None")
+
+    tensor_local = global_to_local(tensor, tensor_layout)
+
+    parallel_net = OneHotExtNet(relu_strategy=(onehot_out_layout,))
+    parallel_output = parallel_net(tensor_local, num_classes)
+
+    parallel_output = local_to_global(parallel_output)
+    assert np.allclose(
+        standalone_output.asnumpy(), parallel_output.asnumpy(), 1e-3, 1e-3
+    )
+
+
+def test_one_hot_ext_3d_data_parallel_int64_4():
+    """
+    Feature: mint.nn.functional.one_hot in python shard.
+    Description: Test with 3D tensor, data parallel on first dim only (indices int64).
+    Expectation: Run success.
+    """
+    ms.set_seed(1)
+    np.random.seed(1)
+    a, b, c, num_classes = 8, 16, 12, 32
+
+    tensor = Tensor(np.random.randint(0, num_classes, size=(a, b, c)).astype(np.int64), ms.int64)
+
+    standalone_net = OneHotExtNet()
+    standalone_output = standalone_net(tensor, num_classes)
+
+    layout = _build_layout()
+    tensor_layout = layout("dp", "None", "None")
+    onehot_out_layout = layout("dp", "None", "None", "None")
+
+    tensor_local = global_to_local(tensor, tensor_layout)
+
+    parallel_net = OneHotExtNet(relu_strategy=(onehot_out_layout,))
+    parallel_output = parallel_net(tensor_local, num_classes)
+
+    parallel_output = local_to_global(parallel_output)
+    assert np.allclose(
+        standalone_output.asnumpy(), parallel_output.asnumpy(), 1e-3, 1e-3
+    )
+
+
+def test_one_hot_ext_auto_depth_skewed_distribution_int64_5():
+    """
+    Feature: mint.nn.functional.one_hot with num_classes=-1 in python shard.
+    Description:
+        Test auto depth inference with skewed data distribution.
+        Only ONE dp shard has the global maximum value to ensure allreduce(max) is required.
+        This avoids the "false positive" where a wrong local-max implementation might pass.
+    Expectation:
+        Run success. Output depth should match global max + 1 on ALL ranks, and global
+        gathered output should match standalone.
+
+    Key point:
+        Under layout("dp") with mesh_shape (2,2,2), cp/mp are replicated.
+        Therefore, all ranks within the same dp shard MUST use identical local indices,
+        otherwise local_to_global() would assemble a different global tensor.
+    """
+    ms.set_seed(1)
+    np.random.seed(1)
+    rank = D.get_rank()
+
+    layout = _build_layout()
+    tensor_layout = layout("dp")
+    onehot_out_layout = layout("dp", "None")
+
+    a = 64
+    expected_depth = 31  # global max is 30
+
+    # Make shard0/shard1 deterministic and identical across all ranks.
+    rng0 = np.random.RandomState(123)
+    rng1 = np.random.RandomState(456)
+
+    shard0 = rng0.randint(0, 8, size=(a,)).astype(np.int64)
+    shard1 = rng1.randint(0, 8, size=(a,)).astype(np.int64)
+
+    # Ensure only dp=1 shard has the global max.
+    shard0[0] = 10
+    shard1[0] = 30
+
+    global_tensor = _make_global_from_dp_shards(shard0, shard1)
+
+    # Standalone reference
+    standalone_net = OneHotExtNet()
+    standalone_out = standalone_net(global_tensor, -1)
+    assert (
+        standalone_out.shape[-1] == expected_depth
+    ), f"Standalone depth should be {expected_depth}, got {standalone_out.shape[-1]}"
+
+    # Parallel: all ranks in the same dp shard use identical local indices.
+    dp_id = _get_dp_id(rank)
+    local_np = shard0.copy() if dp_id == 0 else shard1.copy()
+    tensor_local = DTensor.from_local(Tensor(local_np, ms.int64), tensor_layout)
+
+    parallel_net = OneHotExtNet(relu_strategy=(onehot_out_layout,))
+    parallel_out = parallel_net(tensor_local, -1)
+
+    assert hasattr(
+        parallel_out, "local_shape"
+    ), f"Rank {rank}: output should be DTensor"
+    assert (
+        parallel_out.local_shape[-1] == expected_depth
+    ), f"Rank {rank}: depth should be {expected_depth}, got {parallel_out.local_shape[-1]}"
+
+    parallel_out_global = local_to_global(parallel_out)
+    if rank == 0:
+        assert np.allclose(
+            standalone_out.asnumpy(), parallel_out_global.asnumpy(), 1e-3, 1e-3
+        ), "Parallel output does not match standalone output"
+
+
+def test_one_hot_ext_auto_depth_all_same_local_max_int64_6():
+    """
+    Feature: mint.nn.functional.one_hot with num_classes=-1 in python shard.
+    Description:
+        Test auto depth inference when all dp shards have the same local max (easy case).
+    Expectation:
+        Run success. Output depth matches global max + 1, and global gathered output matches standalone.
+
+    Key point:
+        Still must respect replicate semantics of cp/mp under layout("dp"):
+        ranks within the same dp shard must have identical local indices.
+    """
+    ms.set_seed(2)
+    np.random.seed(2)
+    rank = D.get_rank()
+
+    layout = _build_layout()
+    tensor_layout = layout("dp")
+    onehot_out_layout = layout("dp", "None")
+
+    a = 64
+    expected_depth = 21  # local/global max is 20
+
+    rng = np.random.RandomState(2024)
+    shard0 = rng.randint(0, 15, size=(a,)).astype(np.int64)
+    shard1 = shard0.copy()
+    shard0[0] = 20
+    shard1[0] = 20
+
+    global_tensor = _make_global_from_dp_shards(shard0, shard1)
+
+    standalone_net = OneHotExtNet()
+    standalone_out = standalone_net(global_tensor, -1)
+    assert standalone_out.shape[-1] == expected_depth
+
+    dp_id = _get_dp_id(rank)
+    local_np = shard0.copy() if dp_id == 0 else shard1.copy()
+    tensor_local = DTensor.from_local(Tensor(local_np, ms.int64), tensor_layout)
+
+    parallel_net = OneHotExtNet(relu_strategy=(onehot_out_layout,))
+    parallel_out = parallel_net(tensor_local, -1)
+
+    assert hasattr(
+        parallel_out, "local_shape"
+    ), f"Rank {rank}: output should be DTensor"
+    assert (
+        parallel_out.local_shape[-1] == expected_depth
+    ), f"Rank {rank}: depth should be {expected_depth}, got {parallel_out.local_shape[-1]}"
+
+    parallel_out_global = local_to_global(parallel_out)
+    if rank == 0:
+        assert np.allclose(
+            standalone_out.asnumpy(), parallel_out_global.asnumpy(), 1e-3, 1e-3
+        )
diff --git a/tests/mindspore/st/shard/test_ops_one_hot_ext_shell.py b/tests/mindspore/st/shard/test_ops_one_hot_ext_shell.py
new file mode 100644
index 0000000000000000000000000000000000000000..38e51b4eab3adb0bee9c92d557e365013b4f84cb
--- /dev/null
+++ b/tests/mindspore/st/shard/test_ops_one_hot_ext_shell.py
@@ -0,0 +1,136 @@
+# Copyright 2026 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""parallel_one_hot_ext_shell test"""
+
+from tests.common.mark_utils import arg_mark
+from tests.mindspore.st.utils import msrun_case
+
+
+@arg_mark(
+    plat_marks=["platform_ascend910b"],
+    level_mark="level0",
+    card_mark="allcards",
+    essential_mark="essential",
+)
+def test_one_hot_ext_cell_shard_1():
+    """
+    Feature: OneHotExt operator.
+    Description: Test cell shard in python with 1D tensor, data parallel on tensor dim 0.
+    Expectation: Run success.
+    """
+    glog_v = 2
+    file_name = "one_hot_ext_shard_in_python.py"
+    case_name = "test_one_hot_ext_data_parallel_1d_int64_1"
+    master_port = 12300
+    msrun_case(glog_v, file_name, case_name, master_port)
+
+
+@arg_mark(
+    plat_marks=["platform_ascend910b"],
+    level_mark="level0",
+    card_mark="allcards",
+    essential_mark="essential",
+)
+def test_one_hot_ext_cell_shard_2():
+    """
+    Feature: OneHotExt operator.
+    Description: Test cell shard in python with 2D tensor, data parallel on tensor dim 0.
+    Expectation: Run success.
+    """
+    glog_v = 2
+    file_name = "one_hot_ext_shard_in_python.py"
+    case_name = "test_one_hot_ext_data_parallel_2d_int64_2"
+    master_port = 12301
+    msrun_case(glog_v, file_name, case_name, master_port)
+
+
+@arg_mark(
+    plat_marks=["platform_ascend910b"],
+    level_mark="level0",
+    card_mark="allcards",
+    essential_mark="essential",
+)
+def test_one_hot_ext_cell_shard_3():
+    """
+    Feature: OneHotExt operator.
+    Description: Test cell shard in python with full replication.
+    Expectation: Run success.
+    """
+    glog_v = 2
+    file_name = "one_hot_ext_shard_in_python.py"
+    case_name = "test_one_hot_ext_replicate_all_int64_3"
+    master_port = 12302
+    msrun_case(glog_v, file_name, case_name, master_port)
+
+
+@arg_mark(
+    plat_marks=["platform_ascend910b"],
+    level_mark="level0",
+    card_mark="allcards",
+    essential_mark="essential",
+)
+def test_one_hot_ext_cell_shard_4():
+    """
+    Feature: OneHotExt operator.
+    Description: Test cell shard in python with 3D tensor, data parallel on first dim only.
+    Expectation: Run success.
+    """
+    glog_v = 2
+    file_name = "one_hot_ext_shard_in_python.py"
+    case_name = "test_one_hot_ext_3d_data_parallel_int64_4"
+    master_port = 12303
+    msrun_case(glog_v, file_name, case_name, master_port)
+
+
+@arg_mark(
+    plat_marks=["platform_ascend910b"],
+    level_mark="level0",
+    card_mark="allcards",
+    essential_mark="essential",
+)
+def test_one_hot_ext_cell_shard_5():
+    """
+    Feature: OneHotExt operator with auto depth.
+    Description:
+        Test cell shard in python with 1D tensor, data parallel, num_classes=-1.
+        Only one dp shard has the global maximum to enforce global max reduction.
+    Expectation: Run success with correct global max reduction.
+    """
+    glog_v = 2
+    file_name = "one_hot_ext_shard_in_python.py"
+    case_name = "test_one_hot_ext_auto_depth_skewed_distribution_int64_5"
+    master_port = 12304
+    msrun_case(glog_v, file_name, case_name, master_port)
+
+
+@arg_mark(
+    plat_marks=["platform_ascend910b"],
+    level_mark="level0",
+    card_mark="allcards",
+    essential_mark="essential",
+)
+def test_one_hot_ext_cell_shard_6():
+    """
+    Feature: OneHotExt operator with auto depth.
+    Description:
+        Test cell shard in python with 1D tensor, data parallel, num_classes=-1.
+        All dp shards share the same local max (easy case).
+    Expectation: Run success with correct depth.
+    """
+    glog_v = 2
+    file_name = "one_hot_ext_shard_in_python.py"
+    case_name = "test_one_hot_ext_auto_depth_all_same_local_max_int64_6"
+    master_port = 12305
+    msrun_case(glog_v, file_name, case_name, master_port)
diff --git a/tests/mindspore/ut/parallel_ops_infer/test_parallel_one_hot_ext.py b/tests/mindspore/ut/parallel_ops_infer/test_parallel_one_hot_ext.py
new file mode 100644
index 0000000000000000000000000000000000000000..275c1f8d98e7c241fe0696590e6ff238397c3e83
--- /dev/null
+++ b/tests/mindspore/ut/parallel_ops_infer/test_parallel_one_hot_ext.py
@@ -0,0 +1,201 @@
+# Copyright 2026 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""parallel_one_hot_ext test"""
+
+import pytest
+
+from hyper_parallel import Layout
+from hyper_parallel.core.shard.ops.parallel_one_hot_ext import OneHotExtDistributedOp
+
+
+def run_scenario(scenario_name, indices_layout, expected_map, extra_args):
+    """Infer layout of OneHotExt operator"""
+    print(f"\n{'=' * 80}")
+    print(f"Test OneHotExt, Scenario: {scenario_name}")
+    print("=" * 80)
+
+    op = OneHotExtDistributedOp("OneHotExt")
+    output_layout = op.infer_layout((indices_layout,), extra_args)
+    assert output_layout.to_dict()["tensor_map"] == expected_map, (
+        f"OneHotExt failed in scenario '{scenario_name}'. "
+        f"Expected {expected_map}, got {output_layout.to_dict()['tensor_map']}"
+    )
+
+
+base_mesh_shape = (2, 2, 2)
+base_alias_name = ("dp", "cp", "mp")
+base_rank_list = list(range(8))
+
+
+def _base_layout():
+    return Layout(base_mesh_shape, base_alias_name, base_rank_list)
+
+
+def test_one_hot_ext_1d_axis_minus1_data_parallel_1():
+    """
+    Feature: 1D indices, axis=-1.
+    Description: Indices sharded on dp.
+    Expectation: Output keeps dp sharding and inserts replicated one-hot dim at end.
+    """
+    layout = _base_layout()
+    indices_layout = layout("dp")
+
+    run_scenario(
+        "1. 1D + axis=-1 + DP sharding",
+        indices_layout,
+        expected_map=(2, -1),
+        extra_args=[32, None, None, -1],
+    )
+
+
+def test_one_hot_ext_1d_axis_0_data_parallel_2():
+    """
+    Feature: 1D indices, axis=0.
+    Description: Indices sharded on dp.
+    Expectation: Insert replicated one-hot dim at position 0.
+    """
+    layout = _base_layout()
+    indices_layout = layout("dp")
+
+    run_scenario(
+        "2. 1D + axis=0 + DP sharding",
+        indices_layout,
+        expected_map=(-1, 2),
+        extra_args=[32, None, None, 0],
+    )
+
+
+def test_one_hot_ext_1d_axis_1_data_parallel_3():
+    """
+    Feature: 1D indices, axis=1.
+    Description: axis=1 equals append for rank-1 tensor.
+    Expectation: Same as axis=-1.
+    """
+    layout = _base_layout()
+    indices_layout = layout("dp")
+
+    run_scenario(
+        "3. 1D + axis=1 (append) + DP sharding",
+        indices_layout,
+        expected_map=(2, -1),
+        extra_args=[32, None, None, 1],
+    )
+
+
+def test_one_hot_ext_1d_replicate_all_4():
+    """
+    Feature: 1D indices fully replicated.
+    Description: tensor_map is (-1,).
+    Expectation: Output fully replicated with inserted one-hot dim replicated.
+    """
+    layout = _base_layout()
+    indices_layout = layout("None")
+
+    run_scenario(
+        "4. 1D + replicate all",
+        indices_layout,
+        expected_map=(-1, -1),
+        extra_args=[32, None, None, -1],
+    )
+
+
+def test_one_hot_ext_2d_axis_minus1_data_parallel_5():
+    """
+    Feature: 2D indices, axis=-1.
+    Description: Only data-parallel (dim0 sharded), dim1 replicated.
+    Expectation: Output keeps dp on dim0, inserts replicated one-hot dim at end.
+    """
+    layout = _base_layout()
+    indices_layout = layout("dp", "None")
+
+    run_scenario(
+        "5. 2D + axis=-1 + DP only",
+        indices_layout,
+        expected_map=(2, -1, -1),
+        extra_args=[64, None, None, -1],
+    )
+
+
+def test_one_hot_ext_3d_axis_minus1_data_parallel_6():
+    """
+    Feature: 3D indices, axis=-1.
+    Description: Only data-parallel (dim0 sharded), others replicated.
+    Expectation: Output keeps dp on dim0, inserts replicated one-hot dim at end.
+    """
+    layout = _base_layout()
+    indices_layout = layout("dp", "None", "None")
+
+    run_scenario(
+        "6. 3D + axis=-1 + DP only",
+        indices_layout,
+        expected_map=(2, -1, -1, -1),
+        extra_args=[16, None, None, -1],
+    )
+
+
+def test_one_hot_ext_2d_axis_not_minus1_should_raise_7():
+    """
+    Feature: Multi-dimensional restriction.
+    Description: For rank>1, axis must be -1.
+    Expectation: Raise ValueError.
+    """
+    layout = _base_layout()
+    indices_layout = layout("dp", "None")
+
+    op = OneHotExtDistributedOp("OneHotExt")
+    with pytest.raises(ValueError):
+        _ = op.infer_layout((indices_layout,), [64, None, None, 0])
+
+
+def test_one_hot_ext_2d_not_data_parallel_should_raise_8():
+    """
+    Feature: Multi-dimensional restriction.
+    Description: For rank>1, only dim0 can be sharded (DP only). Sharding dim1 is illegal.
+    Expectation: Raise ValueError.
+    """
+    layout = _base_layout()
+    indices_layout = layout("dp", "cp")
+
+    op = OneHotExtDistributedOp("OneHotExt")
+    with pytest.raises(ValueError):
+        _ = op.infer_layout((indices_layout,), [64, None, None, -1])
+
+
+def test_one_hot_ext_axis_out_of_range_should_raise_9():
+    """
+    Feature: Axis validation.
+    Description: axis must be in [-1, 1] for this implementation.
+    Expectation: Raise ValueError.
+    """
+    layout = _base_layout()
+    indices_layout = layout("dp")
+
+    op = OneHotExtDistributedOp("OneHotExt")
+    with pytest.raises(ValueError):
+        _ = op.infer_layout((indices_layout,), [32, None, None, 2])
+
+
+def test_one_hot_ext_num_classes_invalid_should_raise_10():
+    """
+    Feature: num_classes validation.
+    Description: num_classes must be int and >= -1.
+    Expectation: Raise ValueError when num_classes < -1.
+    """
+    layout = _base_layout()
+    indices_layout = layout("dp")
+
+    op = OneHotExtDistributedOp("OneHotExt")
+    with pytest.raises(ValueError):
+        _ = op.infer_layout((indices_layout,), [-2, None, None, -1])