TPUReshardVariables

public final class TPUReshardVariables

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
create ( Scope scope, Iterable< Operand <?>> vars, Operand <String> newFormatKey, Operand <?> formatStateVar)
Factory method to create a class wrapping a new TPUReshardVariables operation.

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