From 469839cf922504a03704fe46e15bffbc3a80c445 Mon Sep 17 00:00:00 2001
From: Christopher Pruijsen <christopher.pruijsen@gmail.com>
Date: Sat, 12 Sep 2026 01:07:50 +0100
Subject: [PATCH] Fix rank-3 softmax conversion to NeuralNetwork SoftmaxLayer

The NeuralNetwork SoftmaxLayer applies axis=-1 for rank-3 inputs, so
MIL softmax(axis=-3) must lower to SoftmaxND. Rank >= 4 is unchanged.

Fixes #1714
---
 .../converters/mil/backend/nn/op_mapping.py   |  5 ++-
 .../mil/ops/tests/iOS14/test_activation.py    | 34 +++++++++++++++++++
 mlmodel/format/NeuralNetwork.proto            |  5 +--
 3 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/coremltools/converters/mil/backend/nn/op_mapping.py b/coremltools/converters/mil/backend/nn/op_mapping.py
index efc008ab7..99c96f5b7 100644
--- a/coremltools/converters/mil/backend/nn/op_mapping.py
+++ b/coremltools/converters/mil/backend/nn/op_mapping.py
@@ -2472,7 +2472,10 @@ def softplus(const_context, builder, op):
 @register_mil_to_nn_mapping
 def softmax(const_context, builder, op):
     rank = op.x.rank
-    if op.axis.val == -3 or op.axis.val > 0 and op.axis.val == rank - 3:
+    # SoftmaxLayer is specified as axis=-3, but for rank-3 inputs the
+    # NeuralNetwork runtime applies softmax along axis=-1 instead (see issue
+    # 1714). Only emit it when rank >= 4, matching the concat mapping.
+    if rank >= 4 and (op.axis.val == -3 or op.axis.val > 0 and op.axis.val == rank - 3):
         builder.add_softmax(
             name=op.name, input_name=op.x.name, output_name=op.outputs[0].name,
         )
diff --git a/coremltools/converters/mil/mil/ops/tests/iOS14/test_activation.py b/coremltools/converters/mil/mil/ops/tests/iOS14/test_activation.py
index c80f90ddb..2d9513fc3 100644
--- a/coremltools/converters/mil/mil/ops/tests/iOS14/test_activation.py
+++ b/coremltools/converters/mil/mil/ops/tests/iOS14/test_activation.py
@@ -893,6 +893,40 @@ def build(x):
             backend=backend,
         )
 
+    @pytest.mark.parametrize("compute_unit, backend", itertools.product(compute_units, backends))
+    def test_builder_to_backend_rank3_axis_minus_3(self, compute_unit, backend):
+        # Rank-3 softmax(axis=-3) must not use NeuralNetwork SoftmaxLayer
+        # (https://github.com/apple/coremltools/issues/1714): that layer
+        # applies axis=-1 when rank==3, which yields all-ones here.
+        t = np.array([[[0.5]], [[1.5]]], dtype=np.float32)
+        input_placeholders = {"x": mb.placeholder(shape=t.shape)}
+        input_values = {"x": t}
+
+        def build(x):
+            return mb.softmax(x=x, axis=-3)
+
+        expected_output_types = (2, 1, 1, types.fp32)
+        expected_outputs = scipy.special.softmax(t, axis=-3).astype(np.float32)
+        run_compare_builder(
+            build,
+            input_placeholders,
+            input_values,
+            expected_output_types,
+            expected_outputs,
+            compute_unit=compute_unit,
+            backend=backend,
+        )
+
+    def test_neuralnetwork_rank3_axis_minus_3_uses_softmax_nd(self):
+        @mb.program(input_specs=[mb.TensorSpec(shape=(2, 1, 1))])
+        def prog(x):
+            return mb.softmax(x=x, axis=-3, name="y")
+
+        mlmodel = ct.convert(prog, source="milinternal", convert_to="neuralnetwork")
+        layer = mlmodel.get_spec().neuralNetwork.layers[0]
+        assert layer.WhichOneof("layer") == "softmaxND"
+        assert layer.softmaxND.axis == -3
+
     @ssa_fn
     def test_builder_eval(self):
         x_val = np.array([[-1, 2, -3], [4, -5, 6]], dtype=np.float32)
diff --git a/mlmodel/format/NeuralNetwork.proto b/mlmodel/format/NeuralNetwork.proto
index 38dae036b..2f3dacdf1 100644
--- a/mlmodel/format/NeuralNetwork.proto
+++ b/mlmodel/format/NeuralNetwork.proto
@@ -2306,8 +2306,9 @@ message LRNLayerParams {
  * Softmax Normalization Layer
  *
  * A layer that performs softmax normalization.
- * Normalization is applied along axis = -3 or N-3 (where N is the rank of the input)
- * For softmax layer that can operate on any axis, see SoftmaxNDLayer.
+ * For rank >= 4, normalization is applied along axis = -3 or N-3 (where N is the rank of the input).
+ * For rank == 3, the NeuralNetwork runtime applies softmax along axis = -1, not -3.
+ * For a softmax layer that can operate on any axis, see SoftmaxNDLayer.
  *
  *
  * .. code::