tf.raw_ops.XlaSparseDenseMatmul 

tf.raw_ops.XlaSparseDenseMatmul(
    row_ids,
    col_ids,
    values,
    offsets,
    embedding_table,
    max_ids_per_partition,
    max_unique_ids_per_partition,
    input_size,
    name=None
)

row_ids A Tensor of type int32.
col_ids A Tensor of type uint32.
values A Tensor of type float32.
offsets A Tensor of type uint32.
embedding_table A Tensor of type float32.
max_ids_per_partition An int that is >= 0.
max_unique_ids_per_partition An int that is >= 0.
input_size An int that is >= 0.
name A name for the operation (optional).

A tuple of Tensor objects (activations, row_pointers, sorted_embedding_ids, sorted_sample_ids, sorted_gains).
activations A Tensor of type float32.
row_pointers A Tensor of type int32.
sorted_embedding_ids A Tensor of type int32.
sorted_sample_ids A Tensor of type int32.
sorted_gains A Tensor of type float32.