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 | |
| 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.