Skip to content

Make epilogue bits more uniform and useful #628

@t-vi

Description

@t-vi

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 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).

Metadata

Metadata

Assignees

Labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions