Tune in to the first Women in ML Symposium this Tuesday, October 19 at 9am PST Register now

LoadTPUEmbeddingMomentumParameters

public final class LoadTPUEmbeddingMomentumParameters

Load Momentum embedding parameters.

An op that loads optimization parameters into HBM for embedding. Must be preceded by a ConfigureTPUEmbeddingHost op that sets up the correct embedding table configuration. For example, this op is used to install parameters that are loaded from a checkpoint before a training loop is executed.

Nested Classes

class LoadTPUEmbeddingMomentumParameters.Options Optional attributes for LoadTPUEmbeddingMomentumParameters  

Public Methods

static LoadTPUEmbeddingMomentumParameters.Options
config(String config)
static LoadTPUEmbeddingMomentumParameters
create(Scope scope, Operand<Float> parameters, Operand<Float> momenta, Long numShards, Long shardId, Options... options)
Factory method to create a class wrapping a new LoadTPUEmbeddingMomentumParameters operation.
static LoadTPUEmbeddingMomentumParameters.Options
tableId(Long tableId)
static LoadTPUEmbeddingMomentumParameters.Options
tableName(String tableName)

Inherited Methods

Public Methods

public static LoadTPUEmbeddingMomentumParameters.Options config (String config)

public static LoadTPUEmbeddingMomentumParameters create (Scope scope, Operand<Float> parameters, Operand<Float> momenta, Long numShards, Long shardId, Options... options)

Factory method to create a class wrapping a new LoadTPUEmbeddingMomentumParameters operation.

Parameters
scope current scope
parameters Value of parameters used in the Momentum optimization algorithm.
momenta Value of momenta used in the Momentum optimization algorithm.
options carries optional attributes values
Returns
  • a new instance of LoadTPUEmbeddingMomentumParameters

public static LoadTPUEmbeddingMomentumParameters.Options tableId (Long tableId)

public static LoadTPUEmbeddingMomentumParameters.Options tableName (String tableName)