Make PyTorch models up to 40% faster! Thunder is a source to source compiler for PyTorch. It enables using different hardware executors at once; across one or thousands of GPUs.
For making transformations easier and allowing more general return types (e.g. #623) it will be useful to enhance the epilogue uses:
Currently, the prologue return format depends on the presence of an epilogue (it returns either "inputs_to_compute" or"(inputs_to_compute, inputs_to_epilogue)". This leads to early transformation authors needing to consider both cases. If we move to always return also inputs_to_epiloge (an empty tuple if there is no epilogue), this will be more uniform.
Currently, the epilogue does not process the return values. For #623 we might consider constructing the dataclass instance in the epilogue, it would be good to pass the outputs of the compute trace to the epilogue and return the epilogue's return from the function (if there is an epilogue, otherwise, consider it to be an identity function).
For making transformations easier and allowing more general return types (e.g. #623) it will be useful to enhance the epilogue uses:
inputs_to_compute
" or"(inputs_to_compute, inputs_to_epilogue)
". This leads to early transformation authors needing to consider both cases. If we move to always return alsoinputs_to_epiloge
(an empty tuple if there is no epilogue), this will be more uniform.