TableConfig configuration class

[source]

TableConfig class

keras_rs.layers.TableConfig(
    name: str,
    vocabulary_size: int,
    embedding_dim: int,
    initializer: (
        str | keras.src.initializers.initializer.Initializer
    ) = keras.src.initializers.random_initializers.VarianceScaling(mode="fan_out"),
    optimizer: str | keras.src.optimizers.optimizer.Optimizer = "adam",
    combiner: str = "mean",
    placement: str = "auto",
    max_ids_per_partition: int = 256,
    max_unique_ids_per_partition: int = 256,
)

Configuration for one embedding table.

Configures one table for use by one or more keras_rs.layers.FeatureConfig, which in turn is used to configure a keras_rs.layers.DistributedEmbedding.

Attributes

• name: The name of the table. Must be defined.
• vocabulary_size: Size of the table's vocabulary (number of rows).
• embedding_dim: The embedding dimension (width) of the table.
• initializer: The initializer for the embedding weights. If not specified, defaults to truncated_normal_initializer with mean 0.0 and standard deviation 1 / sqrt(embedding_dim).
• optimizer: The optimizer for the embedding table. Only SGD, Adagrad, Adam, and FTRL are supported. Note that not all of the optimizer's parameters are supported. Defaults to Adam.
• combiner: Specifies how to reduce if there are multiple entries in a single row. mean, sqrtn and sum are supported. mean is the default. sqrtn often achieves good accuracy, in particular with bag-of-words columns.
• placement: Where to place the embedding table. "auto", which is the default, means that the table is placed on SparseCore if available, otherwise on the default device where the rest of the model is placed. A value of "sparsecore" means the table will be placed on the SparseCore chips and an error is raised if SparseCore is not available. A value of "default_device" means the table will be placed on the default device where the rest of the model is placed, even if SparseCore is available. The default device for the rest of the model is the TPU's TensorCore on TPUs, otherwise the GPU or CPU.
• max_ids_per_partition: The max number of ids per partition for the table. This is an input data dependent value and is required by the compiler to appropriately allocate memory.
• max_unique_ids_per_partition: The max number of unique ids per partition for the table. This is an input data dependent value and is required by the compiler to appropriately allocate memory.