RetrieveAllTPUEmbeddingParameters

public final class RetrieveAllTPUEmbeddingParameters

An op that retrieves optimization parameters from embedding to host memory.

An op that retrieves optimization parameters from embedding to host memory. Must be preceded by a ConfigureTPUEmbeddingHost op that sets up the correct embedding table configuration. For example, this op is used to retrieve updated parameters before saving a checkpoint. For Adagrad, auxiliary1 will contain the accumulators after running this op. For SGD, all of the auxiliary* values will be empty (0x0 tensors for that table). For FTRL, auxiliary1 will contain the accumulators and auxiliary2 will contain the linear terms. For ADAM, auxiliary1 will contain the momenta and auxiliary2 will contain the velocities.

Public Methods

List<Output<Float>>
auxiliary1()
A list of tensors, one for each embedding table, containing the first auxiliary optimization parameter stored.
List<Output<Float>>
auxiliary2()
A list of tensors, one for each embedding table, containing the second auxiliary optimization parameter stored.
List<Output<Float>>
auxiliary3()
A list of tensors, one for each embedding table, containing the third auxiliary optimization parameter stored.
List<Output<Float>>
auxiliary4()
A list of tensors, one for each embedding table, containing the fourth auxiliary optimization parameter stored.
List<Output<Float>>
auxiliary5()
A list of tensors, one for each embedding table, containing the fifth auxiliary optimization parameter stored.
List<Output<Float>>
auxiliary6()
A list of tensors, one for each embedding table, containing the six auxiliary optimization parameter stored.
List<Output<Float>>
auxiliary7()
A list of tensors, one for each embedding table, containing the seventh auxiliary optimization parameter stored.
static RetrieveAllTPUEmbeddingParameters
create(Scope scope, Long NumTables, String config, Long numShards, Long shardId)
Factory method to create a class wrapping a new RetrieveAllTPUEmbeddingParameters operation.
List<Output<Float>>
parameters()
A list of tensors, one for each embedding table, containing the stored embedding table parameters.

Inherited Methods

Public Methods

public List<Output<Float>> auxiliary1 ()

A list of tensors, one for each embedding table, containing the first auxiliary optimization parameter stored. Elements are present in the list, but have zero size, for unused optimization parameters (based on the algorithm in use for each table).

public List<Output<Float>> auxiliary2 ()

A list of tensors, one for each embedding table, containing the second auxiliary optimization parameter stored. Elements are present in the list, but have zero size, for unused optimization parameters (based on the algorithm in use for each table).

public List<Output<Float>> auxiliary3 ()

A list of tensors, one for each embedding table, containing the third auxiliary optimization parameter stored. Elements are present in the list, but have zero size, for unused optimization parameters (based on the algorithm in use for each table).

public List<Output<Float>> auxiliary4 ()

A list of tensors, one for each embedding table, containing the fourth auxiliary optimization parameter stored. Elements are present in the list, but have zero size, for unused optimization parameters (based on the algorithm in use for each table).

public List<Output<Float>> auxiliary5 ()

A list of tensors, one for each embedding table, containing the fifth auxiliary optimization parameter stored. Elements are present in the list, but have zero size, for unused optimization parameters (based on the algorithm in use for each table).

public List<Output<Float>> auxiliary6 ()

A list of tensors, one for each embedding table, containing the six auxiliary optimization parameter stored. Elements are present in the list, but have zero size, for unused optimization parameters (based on the algorithm in use for each table).

public List<Output<Float>> auxiliary7 ()

A list of tensors, one for each embedding table, containing the seventh auxiliary optimization parameter stored. Elements are present in the list, but have zero size, for unused optimization parameters (based on the algorithm in use for each table).

public static RetrieveAllTPUEmbeddingParameters create (Scope scope, Long NumTables, String config, Long numShards, Long shardId)

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

Parameters
scope current scope
NumTables The number of embedding tables.
config An TPUEmbeddingConfiguration proto describing the table parameters being loaded, serialized to a string.
numShards Number of shards into which the embedding tables are divided.
shardId Identifier of shard for this operation.
Returns
  • a new instance of RetrieveAllTPUEmbeddingParameters

public List<Output<Float>> parameters ()

A list of tensors, one for each embedding table, containing the stored embedding table parameters.