tf.tpu.experimental.embedding.TableConfig

tf.tpu.experimental.embedding.TableConfig(
    vocabulary_size: int,
    dim: int,
    initializer: Optional[Callable[[Any], None]] = None,
    optimizer: Optional[_Optimizer] = None,
    combiner: Text = 'mean',
    name: Optional[Text] = None,
    quantization_config: tf.tpu.experimental.embedding.QuantizationConfig = None,
    layout: Optional[Any] = None
)

table_config_one = tf.tpu.experimental.embedding.TableConfig(
    vocabulary_size=...,
    dim=...)
table_config_two = tf.tpu.experimental.embedding.TableConfig(
    vocabulary_size=...,
    dim=...)
feature_config = {
    'feature_one': tf.tpu.experimental.embedding.FeatureConfig(
        table=table_config_one),
    'feature_two': tf.tpu.experimental.embedding.FeatureConfig(
        table=table_config_one),
    'feature_three': tf.tpu.experimental.embedding.FeatureConfig(
        table=table_config_two)}
embedding = tf.tpu.experimental.embedding.TPUEmbedding(
    feature_config=feature_config,
    batch_size=...
    optimizer=tf.tpu.experimental.embedding.Adam(0.1))