Implement transformation for TensorFlow 2 Map Function (aka tf.map_fn)

[rkazants]

Root cause analysis: In several TF2 OD models (MobileNet SSD V2, CenterNet) we have a preprocessor block based on tf.map_fn primitive that slices multiple inputs, transform them, and concatenate multiple intermediate outputs at each iteration of While Loop. Currently, the MO tool does not support such block.

Solution: Extend existing transformations for TF2 Keras RNN operations that are particular case of tf.map_fn when it has just one sliced input and one concatenated output. The proposed extension detects sub-graphs (from the main-graph and body-graph) for slicing and concatenation and utilizes axis attribute of Loop operation for such operations.

Ticket: 47280

Code:

• Comments
• Code style (PEP8)
• Transformation generates reshape-able IR
• Transformation preserves original framework node names

Validation:

• Unit tests - N/A due to complexity of such unit-tests implementation
• Framework operation tests
• Transformation tests - N/A due to complexity of such unit-tests implementation. Now it covers by TF2 Keras RNN layer tests
• Model Optimizer IR Reader check

Documentation:

• Supported frameworks operations list - N/A
• Guide on how to convert the public model - N/A
• User guide update - N/A

[achetver]

Looks good to me

[lazarevevgeny]

There are a number of minor comments, but in general looks good. However, I would really want to see a new layer test covering the added functionality:

1. A test case where a model has one input which is consumed with the TF While operation which produces 2 outputs.
2. A test case where a model has two inputs which are consumed with the TF While operation which produces 2 outputs.

If one of these test cases is already covered with the existing Keras RNN-like tests that is fine. But the case which was not previously supported should be added to the pre-commit layer tests.

[rkazants]

There are a number of minor comments, but in general looks good. However, I would really want to see a new layer test covering the added functionality:

1. A test case where a model has one input which is consumed with the TF While operation which produces 2 outputs.
2. A test case where a model has two inputs which are consumed with the TF While operation which produces 2 outputs.

If one of these test cases is already covered with the existing Keras RNN-like tests that is fine. But the case which was not previously supported should be added to the pre-commit layer tests.

Thanks for your point. I will add it.