Op that reshards on-device TPU variables to specified state.
Op that reshards on-device TPU variables to specified state. Internal use only.
The sharding state is represented as the key of the compilation that generated the sharding/unsharding programs along with the main program. new_format_key specifies the desired state, and format_state_var is the current state of the variables.
Public Methods
static TPUReshardVariables |
Inherited Methods
Public Methods
public static TPUReshardVariables create ( Scope scope, Iterable< Operand <?>> vars, Operand <String> newFormatKey, Operand <?> formatStateVar)
Factory method to create a class wrapping a new TPUReshardVariables operation.
Parameters
scope | current scope |
---|
Returns
- a new instance of TPUReshardVariables