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