einops.rearrange does not permit to export model

[doloresgarcia]

A simple dimension change gives error when exporting model with torch.onnx.dynamo_export a = rearrange(a,"... c x -> ... (c x)", ) Expected result: torch.Size([1, 8, 565, 4, 7]) torch.Size([1, 8, 565, 28])

torch '2.2.2+cpu' onnxscript '0.1.0.dev20240925' onnx '1.16.2'

[justinchuby]

Can you use torch-nightly, export with the torch.onnx.export(…, dynamo=True, report=True, verify=True) call and share the markdown report?

[doloresgarcia]

Hi @justinchuby, thank you for your help! Here is the markdown. The terminal shows this:

torch._dynamo.exc.TorchRuntimeError: Failed running call_function <function rearrange at 0x7f41f0859cf0>(*(FakeTensor(..., size=(1, 8, s0, 4, 7), grad_fn=<IndexSelectBackward0>), '... c x -> ... (c x)'), **{}):
unhashable type: non-nested SymInt

onnx_export_2024-09-26_10-08-04-071036_pt_export.md

If I move to einops 0.8.0 and do not ask for dynamic shapes it does work but when moving to dynamic axis the result shows:

torch.fx.experimental.symbolic_shapes.ConstraintViolationError: Constraints violated (inputs_dim_0)! For more information, run with TORCH_LOGS="+dynamic".
E0926 14:05:23.241000 23175 site-packages/torch/_guards.py:282] [0/0]   - Not all values of inputs_dim_0 = L['inputs'].size()[0] in the specified range satisfy the generated guard Ne(L['inputs'].size()[0], 9223372036854775807).

[Uploading onnx_export_2024-09-26_14-05-12-652712_pt_export.md…]()

[justinchuby]

Looks like einops 0.8 has fixed some issues then. The second report is not uploaded. Could you share that?

[doloresgarcia]

the problem seems to be solved by using torch '2.2.4+cpu' and einops 0.8.0. It is not clear where the guard error comes from but it does not seem to be related to the einops.rearrange and it worked after modifying statements of this kind: x[:,0], I will close the issue. Thank you @justinchuby