[Grappler] Add RemoveStackStridedSliceSameAxis optimizer.

// Replace operations of the form:
//    x = stack((a_0, a_1, ..., a_{n-1}), axis=k)[:,...,i,...]
// with
//    a_i
// when the strided slice index `i` is applied in the k'th axis.
//
// Similarly, replace operations of the form:
//    x = stack((a_0, a_1, ..., a_{n-1}), axis=k)[:,...,i:i+1,...]
// with
//    expand_dims(a_i, axis=k)
//

PiperOrigin-RevId: 216535346

1 files changed, 3 insertions, 0 deletions

diff --git a/tensorflow/core/grappler/optimizers/graph_optimizer_stage_test.cc b/tensorflow/core/grappler/optimizers/graph_optimizer_stage_test.cc
index 34f28c7c27..799c40c67b 100644
--- a/tensorflow/core/grappler/optimizers/graph_optimizer_stage_test.cc
+++ b/tensorflow/core/grappler/optimizers/graph_optimizer_stage_test.cc
@@ -61,6 +61,7 @@ TEST_F(GraphOptimizerStageTest, OptimizedNodeName) {
                             /*optimized_graph*/ nullptr,
                             /*graph_properties*/ nullptr,
                             /*node_name*/ nullptr,
+                            /*feed_nodes*/ nullptr,
                             /*opt_level*/ RewriterConfig::ON);
   FakeOptimizerStage stage("my_opt", "my_stg", ctx);
 
@@ -97,6 +98,7 @@ TEST_F(GraphOptimizerStageTest, GetInputNodeAndProperties) {
                             /*optimized_graph*/ &item.graph,
                             /*graph_properties*/ &properties,
                             /*node_name*/ &node_map,
+                            /*feed_nodes*/ nullptr,
                             /*opt_level*/ RewriterConfig::ON);
   FakeOptimizerStage stage("my_opt", "my_stg", ctx);
 
@@ -137,6 +139,7 @@ TEST_F(GraphOptimizerStageTest, AddNodes) {
                             /*optimized_graph*/ &item.graph,
                             /*graph_properties*/ &properties,
                             /*node_name*/ &node_map,
+                            /*feed_nodes*/ nullptr,
                             /*opt_level*/ RewriterConfig::ON);
   FakeOptimizerStage stage("my_opt", "my_stg", ctx);