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