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

org.tensorflow.op.PrimitiveOp
final boolean
equals(Object obj)
final int
hashCode()
Operation
op()
Returns the underlying Operation
final String
toString()
boolean
equals(Object arg0)
final Class<?>
getClass()
int
hashCode()
final void
notify()
final void
notifyAll()
String
toString()
final void
wait(long arg0, int arg1)
final void
wait(long arg0)
final void
wait()

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