torch.jit.frontend.NotSupportedError: keyword-arg expansion is not supported: for dgl.nn.HeteroGraphConv()

[sabadaftari]

🚀 The feature, motivation and pitch

Hello,

I get this error once I call the following line after my model is trained.

model_scripted = torch.jit.script(model) # Export to TorchScript
model_scripted.save('model_scripted.pt') # Save

my model has a HeteroGraphConv layer that uses dictionary-based arguments (two node types). However, TorchScript has limitations compared to dynamic Python execution, and keyword argument expansion is one of the features not supported in TorchScript.

Not sure if it's a bug or a feature to be developed, but NotSupportedError made me think it's a feature request.

Thanks!

Alternatives

No response

Additional context

torch.jit.frontend.NotSupportedError: keyword-arg expansion is not supported:
  File "../opt/anaconda3/lib/python3.9/site-packages/dgl/nn/pytorch/hetero.py", line 178
                    (src_inputs[stype], dst_inputs[dtype]),
                    *mod_args.get(etype, ()),
                    **mod_kwargs.get(etype, {}))
                      ~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
                outputs[dtype].append(dstdata)
        else:

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel

[jgong5]

Yes, it is not supported. Are you able to compile the model with torch.compile? torch.jit is in maintenance mode and I don't think we will add new features to it...

[sabadaftari]

I think torch.compile() does not save the model. I want to save my model weights with saving the classes at the same time, so that when we send it for production, they don't have to define the models classes. That's why I want to save it in TorchScript format.

[jgong5]

I think torch.compile() does not save the model. I want to save my model weights with saving the classes at the same time, so that when we send it for production, they don't have to define the models classes. That's why I want to save it in TorchScript format.

Can torch.jit.trace work from your side? Check the UT here for passing dict: https://github.com/pytorch/pytorch/pull/81623/files#diff-c54c3bae2ea206c4de5b90adf7f2d3ec47df6818451b8dd2b3c36fe54adc10d8R3030