nn.Module.forward signature with **kwargs

[prokotg]

Hi there!

The other day I was working on improving utils.checkpoint module and I found out there is no way right now to pass **kwargs to checkpointed function. This is probably in order to by compliant with the nn.Module.forward method.

Now, I could not think about any argument against having **kwargs in forward method and many people overload this method with **kwargs because it is nice and lit :) I was about to change it and create pull request but before I do it I decided to go smart and ask you guys if there is some obvious reason why we should not do it. Otherwise, I am more than happy to start working on it.

Thanks :- )

[fmassa]

Adding a **kwargs argument to the base forward method won't change anything in the functionality I believe, given that __call__ already has thr kwargs and passes it to forward, see

pytorch/torch/nn/modules/module.py

Line 529 in b2f6e2b

def __call__(self, *input, **kwargs):

This being said, I don't think there is any problem exposing the kwargs in the base forward

[ssnl]

There is nothing that prevents it but you gain nothing from this change as well because forward is always overriden.

[prokotg]

True, no gain in functionality, just a change to be consistent I guess. I am going to try to change checkpoint module to take into account **kwargs though.

Thank you :)

[saareliad]

Currently it could be solved if programmers write forward functions in a way which checkpoint will be able to pass *args to it (in order) instead using kwargs.

OR

Pass a dict function_kwargs to checkpoint, which will end up calling:

CheckpointFunction.apply(function, preserve, *args, **function_kwargs)
(and change CheckpointFunction accordingly)
instead of
CheckpointFunction.apply(function, preserve, *args)

For example, function_kwargs could be added as a kwarg to the checkpoint function.

IMO should be easy to implement in a backward compatible way.

[prokotg]

@saareliad right, that's what the documentation says. The forward function should know what to do with **args so we have to pass it in some order and indeed you have to know what to do with input tuple. However, your forward method might require, for example, some parameters and you can have multiple of them and passing them without keyword could get messy.

Besides, current solution forces developers to rewrite their current, key-worded model to be keyword-args-less and I think since pytorch is aiming to be python-first, this should not take place.

I followed the data-flow path and basically started to change some code to incorporate kwargs
here

pytorch/torch/csrc/autograd/python_function.cpp

Line 1056 in b2f6e2b

{(char*)"apply", (PyCFunction)THPFunction_apply, METH_CLASS | METH_VARARGS, nullptr},

and in the function itself to unpack the **kwargs dict

pytorch/torch/csrc/autograd/python_function.cpp

Line 675 in b2f6e2b

PyObject *THPFunction_apply(PyObject *cls, PyObject *inputs)

and I stumbled upon #16940