Fix layernorm and softmax axis after upstream (microsoft#17255)

### Fix layernorm and softmax axis after upstream

For Gather (the slicing is a scalar), the output rank is small than its
inputs.

When we upstream this kind of Gather before softmax or layernorm, we
should also update the axis attribute.
Otherwise, the axis might be out-of-date and incorrect for the updated
rank.

```
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_fallback.py", line 157, in handle_exception
    raise exception
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_training_manager.py", line 280, in forward
    self._build_graph(graph_transformer_config)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_logger.py", line 158, in wrapper
    result = func(graph_execution_manager, *args, **kwargs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_logger.py", line 273, in wrapper
    result = func(graph_execution_manager, *args, **kwargs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_training_manager.py", line 361, in _build_graph
    super()._build_graph(graph_transformer_config)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_graph_execution_manager.py", line 184, in _build_graph
    self._graph_builder.build(config)
RuntimeError: /onnxruntime/orttraining/orttraining/python/orttraining_pybind_state.cc:823 onnxruntime::python::addObjectMethodsForTraining(pybind11::module&, onnxruntime::python::ExecutionProviderRegistrationFn)::<lambda(onnxruntime::training::OrtModuleGraphBuilder*, const onnxruntime::training::TrainingGraphTransformerConfiguration&)> [ONNXRuntimeError] : 1 : FAIL : Node (Softmax_2904) Op (Softmax) [ShapeInferenceError] 'axis' must be in [-3 , 2]. Its actual value is: 3
```

onnxruntime/core/optimizer/compute_optimizer/upstream_gather_actors.cc

@@ -462,6 +462,27 @@ bool LayerNormalizationGatherActor::PreCheck(const Graph& /* graph */,
   return true;
 }
 
+bool LayerNormalizationGatherActor::PostProcess(Graph& /*graph*/, Node& current_node,
+                                                const SliceInfo& info_without_node,
+                                                const logging::Logger& /*logger*/,
+                                                const std::unordered_map<int, int>& /*propagate_input_indices*/,
+                                                const std::unordered_map<int, std::vector<DimCompare>>&
+                                                /*all_input_cmp_rets*/,
+                                                const std::unordered_map<int, SliceInfo>& /*new_gather_infos*/) {
+  // Update LayerNormalization's axis attribute if it is scalar slice.
+  if (info_without_node.is_scalar_slice) {
+    auto axis = static_cast<int64_t>(current_node.GetAttributes().at("axis").i());
+    auto original_ln_input_rank = info_without_node.input_rank;
+    axis = axis < 0 ? axis + original_ln_input_rank : axis;
+    auto new_axis = axis - 1;
+
+    auto& attributes = current_node.GetMutableAttributes();
+    attributes["axis"] = ONNX_NAMESPACE::MakeAttribute("axis", static_cast<int64_t>(new_axis));
+  }
+
+  return true;
+}
+
 bool SoftmaxGatherActor::PreCheck(const Graph& graph, const Node& current_node, const SliceInfo& info,
                                   const logging::Logger& logger,
                                   std::unordered_map<int, int>& propagate_input_indices,
@@ -479,6 +500,28 @@ bool SoftmaxGatherActor::PreCheck(const Graph& graph, const Node& current_node,
                                                     propagate_input_indices, all_input_cmp_rets, shape_update_func);
 }
 
+bool SoftmaxGatherActor::PostProcess(Graph& graph, Node& current_node, const SliceInfo& info_without_node,
+                                     const logging::Logger& logger,
+                                     const std::unordered_map<int, int>& propagate_input_indices,
+                                     const std::unordered_map<int, std::vector<DimCompare>>& all_input_cmp_rets,
+                                     const std::unordered_map<int, SliceInfo>& new_gather_infos) {
+  SimplePointwiseGatherActor<true>::PostProcess(graph, current_node, info_without_node, logger,
+                                                propagate_input_indices, all_input_cmp_rets, new_gather_infos);
+
+  // Update Softmax's axis attribute if it is scalar slice.
+  if (info_without_node.is_scalar_slice) {
+    auto axis = static_cast<int64_t>(current_node.GetAttributes().at("axis").i());
+    auto original_ln_input_rank = info_without_node.input_rank;
+    axis = axis < 0 ? axis + original_ln_input_rank : axis;
+    auto new_axis = axis - 1;
+
+    auto& attributes = current_node.GetMutableAttributes();
+    attributes["axis"] = ONNX_NAMESPACE::MakeAttribute("axis", static_cast<int64_t>(new_axis));
+  }
+
+  return true;
+}
+
 bool ReshapeGatherActor::PreCheck(const Graph& graph, const Node& current_node, const SliceInfo& info,
                                   const logging::Logger& logger,
                                   std::unordered_map<int, int>& propagate_input_indices,
@@ -566,6 +609,11 @@ bool ReshapeGatherActor::PreCheck(const Graph& graph, const Node& current_node,
     return true;
   }
 
+  LOG_DEBUG_INFO(logger, "Skip handle the Reshape, new_shape_const_values[info.non_negative_axis]:" +
+                             std::to_string(new_shape_const_values[info.non_negative_axis]) +
+                             ", info.output_dim_on_axis.has_dim_value(): " +
+                             std::to_string(info.output_dim_on_axis.has_dim_value()) + ".");
+
   return false;
 }
 
@@ -604,11 +652,12 @@ bool ReshapeGatherActor::PostProcess(
     return true;
   }
 
-  // If it selected shape is a dim value, we can update the shape tensor directory.
+  // If the selected shape is a dim value, we can update the shape tensor directory.
   if (info_without_node.output_dim_on_axis.has_dim_value()) {
     new_shape_const_values[slice_axis] = info_without_node.output_dim_on_axis.dim_value();
     auto new_shape_arg =
-        CreateInitializerFromVector(graph, {static_cast<int64_t>(new_shape_const_values.size())}, new_shape_const_values,
+        CreateInitializerFromVector(graph, {static_cast<int64_t>(new_shape_const_values.size())},
+                                    new_shape_const_values,
                                     graph.GenerateNodeArgName(current_node.MutableInputDefs()[1]->Name()));
     graph_utils::ReplaceNodeInput(current_node, 1, *new_shape_arg);
     return true;

onnxruntime/core/optimizer/compute_optimizer/upstream_gather_actors.h

@@ -189,7 +189,7 @@ class LayerNormalizationGatherActor : public UpStreamGatherOperatorActorBase {
                    const logging::Logger& /* logger */,
                    const std::unordered_map<int, int>& /* propagate_input_indices */,
                    const std::unordered_map<int, std::vector<DimCompare>>& /* all_input_cmp_rets */,
-                   const std::unordered_map<int, SliceInfo>& /* new_gather_infos */) override { return true; }
+                   const std::unordered_map<int, SliceInfo>& /* new_gather_infos */) override;
 };
 
 class SoftmaxGatherActor : public SimplePointwiseGatherActor<true> {
@@ -202,6 +202,12 @@ class SoftmaxGatherActor : public SimplePointwiseGatherActor<true> {
                 std::unordered_map<int, int>& propagate_input_indices,
                 std::unordered_map<int, std::vector<DimCompare>>& all_input_cmp_rets,
                 std::function<void(Node& node)>& shape_update_func) override;
+
+  bool PostProcess(Graph& /* graph */, Node& /* current_node */, const SliceInfo& /* info_without_node */,
+                   const logging::Logger& /* logger */,
+                   const std::unordered_map<int, int>& /* propagate_input_indices */,
+                   const std::unordered_map<int, std::vector<DimCompare>>& /* all_input_cmp_rets */,
+                   const std::unordered_map<int, SliceInfo>& /* new_gather_infos */) override;
 };
 
 class ReshapeGatherActor : public UpStreamGatherOperatorActorBase {