Add Softmax variant for rank-3 with 2 reduction axes (#1360)

src/ngraph/runtime/cpu/builder/softmax.cpp

@@ -98,6 +98,19 @@ namespace ngraph
                         };
                         functors.emplace_back(functor);
                     }
+                    else if (arg_shape.size() == 3 && axes.size() == 2)
+                    {
+                        std::function<decltype(runtime::cpu::kernel::softmax_3d_2rd<float>)> kernel;
+
+                        SELECT_KERNEL(kernel,
+                                      args[0].get_element_type(),
+                                      runtime::cpu::kernel::softmax_3d_2rd);
+
+                        auto functor = [&, kernel, arg_shape, axes](CPURuntimeContext* ctx) {
+                            kernel(arg_tensor, out_tensor, arg_shape, axes);
+                        };
+                        functors.emplace_back(functor);
+                    }
                     else
                     {
                         throw ngraph_error("Unsupported Softmax");

src/ngraph/runtime/cpu/kernel/softmax.hpp

@@ -106,6 +106,15 @@ namespace ngraph
                 {
                     softmax<ElementType, Rank, 1>(input, output, input_shape, softmax_axes);
                 }
+
+                template <typename ElementType>
+                void softmax_3d_2rd(void* input,
+                                    void* output,
+                                    const Shape& input_shape,
+                                    const AxisSet& softmax_axes)
+                {
+                    softmax<ElementType, 3, 2>(input, output, input_shape, softmax_axes);
+                }
             }
         }
     }