TensorDimensionMismatchException on Matmul

[JeremyMorlier]

When considering a Matmul node that uses the output of a conv layer followed by Transpose and Reshape nodes (as done in training) as the B part of it, Stream returns an error that I can't figure out.

The error

Traceback (most recent call last):
  File "PATH/stream/stream/stages/generation/tiled_workload_generation.py", line 650, in get_inter_edges_numpy
    inter_edges = self.get_inter_edges_tensor_based(tensor, final_tensor)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "PATH/stream/stream/stages/generation/tiled_workload_generation.py", line 715, in get_inter_edges_tensor_based
    raise TensorDimensionMismatchException("Arrays to construct inter-layer edges must be equal shape.")
stream.stages.generation.tiled_workload_generation.TensorDimensionMismatchException: Arrays to construct inter-layer edges must be equal shape.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "PATH/stream/stream/stages/generation/tiled_workload_generation.py", line 619, in get_final_tensor_alt_operand
    alt_operand = next(op for op in final_node.input_operand_source if op != dependent_operand)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
StopIteration

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "PATH/stream/stream/stages/generation/tiled_workload_generation.py", line 654, in get_inter_edges_numpy
    final_tensor = get_final_tensor_alt_operand()
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "PATH/stream/stream/stages/generation/tiled_workload_generation.py", line 622, in get_final_tensor_alt_operand
    raise TensorDimensionMismatchException
stream.stages.generation.tiled_workload_generation.TensorDimensionMismatchException

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "PATH/stream/stream/stages/generation/tiled_workload_generation.py", line 660, in get_inter_edges_numpy
    tensor = get_shape_inferred_propagated_tensor(tensor, final_tensor)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "PATH/stream/stream/stages/generation/tiled_workload_generation.py", line 630, in get_shape_inferred_propagated_tensor
    raise TensorDimensionMismatchException(
stream.stages.generation.tiled_workload_generation.TensorDimensionMismatchException: This function only solves the case of errors due to constant shapes in ConcatNode

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "PATH/stream/stream/stages/generation/tiled_workload_generation.py", line 619, in get_final_tensor_alt_operand
    alt_operand = next(op for op in final_node.input_operand_source if op != dependent_operand)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
StopIteration

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "PATH/error_matmul.py", line 43, in <module>
    scme = optimize_allocation_ga(
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "PATH/stream/stream/api.py", line 110, in optimize_allocation_ga
    answers = mainstage.run()
              ^^^^^^^^^^^^^^^
  File "PATH/stream/stream/stages/stage.py", line 62, in run
    for cme, extra_info in self.list_of_callables[0](self.list_of_callables[1:], **self.kwargs).run():
  File "PATH/stream/stream/stages/parsing/accelerator_parser.py", line 29, in run
    for cme, extra_info in sub_stage.run():
  File "PATH/stream/stream/stages/parsing/onnx_model_parser.py", line 42, in run
    for cme, extra_info in sub_stage.run():
  File "PATH/stream/stream/stages/generation/layer_stacks_generation.py", line 72, in run
    for cme, extra_info in sub_stage.run():
  File "PATH/stream/stream/stages/generation/tiling_generation.py", line 52, in run
    for cme, extra_info in sub_stage.run():
  File "PATH/stream/stream/stages/generation/tiled_workload_generation.py", line 97, in run
    inter_edges = self.get_inter_edges_numpy(producer, consumer)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "PATH/stream/stream/stages/generation/tiled_workload_generation.py", line 665, in get_inter_edges_numpy
    final_tensor = get_final_tensor_alt_operand()
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "PATH/stream/stream/stages/generation/tiled_workload_generation.py", line 622, in get_final_tensor_alt_operand
    raise TensorDimensionMismatchException
stream.stages.generation.tiled_workload_generation.TensorDimensionMismatchException

Some code to reproduce the error :

import torch
import torch.nn as nn
import torch.nn.functional as F

import onnx
from onnx import shape_inference
from onnxruntime.training import artifacts

from stream.api import optimize_allocation_ga

class Reproducer(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(10, 20, (3, 3))
        self.linear2 = nn.Linear(20, 10)
    def forward(self, x):
        x = self.conv1(x)
        x = torch.permute(x, (0, 2, 3, 1))
        x = torch.reshape(x, (x.shape[0], x.shape[1] * x.shape[2], x.shape[3]))
        print(x.shape)
        x = torch.matmul(torch.ones((x.shape[0], 10, x.shape[1])), x)
        return x
    
if __name__ == "__main__":
    folder = "onnx/error_matmul"
    onnx_path = f"{folder}/test.onnx"
    infered_path =  f"{folder}/inferred.onnx"
    train_onnx_path = f"{folder}/training_model.onnx"
    inferred_train_onnx_path = f"{folder}/infered_training_model.onnx"
    inferred_train_onnx_path2 = f"{folder}/infered_training_model2.onnx"
    soc_path = "stream/stream/inputs/examples/hardware/tpu_like_quad_core.yaml"
    mapping_path = "stream/stream/inputs/examples/mapping/tpu_like_quad_core.yaml"
    output_path = "output/result"
    mode = "fused"

    layer_stacks = [tuple(range(0, 11)), tuple(range(11, 22))] + list((i,) for i in range(22, 49))

    model = Reproducer()
    torch_input = torch.randn(4, 10, 32, 32)
    torch.onnx.export(model, torch_input, onnx_path, opset_version=13)
    inferred_model = shape_inference.infer_shapes_path(onnx_path, infered_path)

    scme = optimize_allocation_ga(
        hardware=soc_path,
        workload=infered_path,
        mapping=mapping_path,
        mode=mode,
        layer_stacks=layer_stacks,
        nb_ga_generations=4,
        nb_ga_individuals=4,
        experiment_id=id,
        output_path=output_path,
        skip_if_exists=False,
    )

the network architecture