tf-encrypted / moose

Secure distributed dataflow framework for encrypted machine learning and data processing
Apache License 2.0
58 stars 15 forks source link

Impute placements in AbstractComputation #1133

Closed jvmncs closed 2 years ago

jvmncs commented 2 years ago

Closes #1112

Allows for imputing new placements into AbstractComputations during tracing via a role_map, which is a map of placement names to placement names. I considered making this a mapping of PlacementExpression->PlacementExpression or Placement->Placement, but did not want to require having these objects in scope when building the role_map dict.

The role_map can be provided either to the pm.computation decorator directly, or to an existing AbstractComputation class via the with_role_map method.