torch-mlir icon indicating copy to clipboard operation
torch-mlir copied to clipboard

[FxImporter]: support parsing nodes which return list

Open Vremold opened this issue 1 year ago • 7 comments

As described in title. The context is that I'm using torch.export + torch_mlir.fx_importer to trace a Mixtral MoE layer. It has aten.unbind.int op, which returns a list of tensor as the output.

Vremold avatar Mar 16 '24 10:03 Vremold

Thanks for the contribution. If you wouldn't mind, could you add a test? I know the testing in tree is thin for this (we have a bunch of tests out of tree that have not been ported yet), but would like to not keep adding features without tests for structural things.

Existing tests are here: https://github.com/llvm/torch-mlir/tree/main/test/python/fx_importer

Could you create a new op_forms.py and add a case. It should follow how basic_test.py does it. Checks don't need to be super detailed.

Thanks for the kind advice. A simple test case for list_return_test.py which contains only one case of torch.unbind.int op is added.

Vremold avatar Mar 18 '24 16:03 Vremold

Thank you for the contribution. Since this stuff gets a bit twisty, do you mind if I take a bit of time and poke at this so that I understand?

OK.

Vremold avatar Mar 22 '24 01:03 Vremold