Closed mortendahl closed 2 years ago
Since computations are rewritten and operations renamed, this might require adding an attribute on output operations that can be used to reconstruct eg the order specified in PyMoose.
I think the "right" way to preserve output ordering would be to preserve the op names from the ASTTracer. this could be achieved if we were to implement name scoping for ops, e.g. so that lowered ops could try to maintain a record of the higher-level ops they came from via their name. this would be consistent with what we saw in TensorFlow/TF Encrypted. an example for replicated addition might be logical level add_0
turning into a subgraph of ops with names like add_0/share_0/ring_add_0
, add_0/replicated_add_0/ring_add_0
, etc. (these examples are not meant to be correct, just for illustrative purposes)
but I admit name scoping might be an ambitious goal for right now, so think your suggestion for an explicit pm.output
makes sense
A computation like the following currently results in an output op named
op_3
:Is there any way we could provide better names for outputs, or at least ensure ordering?
For instance:
or