HFTracer.trace can now take callables and torch.nn.Module

[michaelbenayoun]

What does this PR do?

This PR enables to use the HFTracer "meta-tracing" features to trace any Python callable / torch.nn.Module.

For transformers.PreTrainedModels, the method HFTracer._generate_dummy_inputs already takes care of creating the original dummy inputs needed to handle data-dependent control-flow in the forward pass.

Now, the user can specify dummy_inputs directly to the HFTracer.trace method in order to be able to trace other things than transformers.PreTrainedModels. This is useful for pattern matching for instance.

This becomes possible:

def f(x, y, z=None):
    temp =  x * y
    if z is not None:
        temp += z
    return temp

traced_f = HFTracer().trace(f, dummy_inputs={"x": torch.rand(1, 2), "y": torch.rand(1, 2)})

By default, if dummy_inputs is specified, every argument to root that is not in dummy_inputs will be considered a concrete arg (and thus added to concrete_args). You can disable that by setting infer_concrete_args_from_dummy_inputs to False. This is useful if want to provide custom dummy inputs for some inputs, while still keeping the HFTracer._generate_dummy_inputs doing the work for other inputs (provided that root is a transformers.PreTrainedModel since only this case is supported for automatic dummy inputs generation).

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[thomasw21]

Nice! Some small comment!

[thomasw21]

Suggested change

	infer_concrete_args_from_dummy_inputs (`bool`, *optional*, defaults to `True`):
	infer_concrete_args_from_dummy_inputs (`bool`, defaults to `True`):

[sgugger]

No, the argument is optional since it has a default (the user does not have to provide it). Please read the writing documentation guide.

[thomasw21]

I was unaware of this. I mistakenly thought that bool, *optional* means Optional[bool] in typing nomenclature. This seems weird to me as all defaults become optional now? Anyway perhaps a chat we can have someplace else.

[sgugger]

All arguments that have a default are optionals, yes. That is the definition of an optional argument, an argument you do not need to provide. It's not because the typing module decide to (badly) reuse that word for something else that this will change.

[thomasw21]

What happens if the parameter doesn't have default?

[michaelbenayoun]

Then tracing will most likely fail afterwards, added a check to fail early, as you suggested.

[sgugger]

Thanks for adding this!

[sgugger]

No, the argument is optional since it has a default (the user does not have to provide it). Please read the writing documentation guide.