-
Notifications
You must be signed in to change notification settings - Fork 112
Make epilogue bits more uniform and useful #628
Copy link
Copy link
Closed
Labels
Description
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 alsoinputs_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 support returning dataclasses with tensors #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).
Reactions are currently unavailable