I can confirm the same issue exists even when converting models without state. I do not think this check makes sense. For example, my model has two inputs. The first input J has 3 shapes. The second input K has two shapes, and the correspondence goes like, J[0] <-> K[0] J[1] <-> K[0] J[2] <-> K[1] and there is currently no way to to this I think