dynamic_embedding does not support tf.RaggedTensor as the input to look up

[cockroachzl]

Describe the feature and the current behavior/state.
First of all, thank you all for making this awesome project!

Right now dynamic embedding seems does not support embedding lookup with ragged tensor. For example, this code snippet for dynamic embedding lookup works well with regular tf.tensor

movie_embeddings = tfra.dynamic_embedding.get_variable(
    name="movie_dynamic_embeddings",
    dim=16,
    initializer=tf.keras.initializers.RandomNormal(-1, 1))

movie_embeddings_shadow = tfra.dynamic_embedding.shadow_ops.ShadowVariable(
    movie_embeddings,
    name='movie_dynamic_embeddings_shadow',
    max_norm=None,
    trainable=True)

movie_ids = tf.random.uniform(shape=[8, 4], minval=0, maxval=15, dtype=tf.int64)

movie_id_weights = tfra.dynamic_embedding.shadow_ops.embedding_lookup(movie_embeddings_shadow, movie_ids)

However, if movie_ids is a ragged tensor, like

ragged_movie_ids = tf.ragged.constant([[1, 2], [3], [4, 5, 6]])
movie_id_weights = tfra.dynamic_embedding.shadow_ops.embedding_lookup(movie_embeddings_shadow, ragged_movie_ids)

doing the embedding lookup failed with the following stack trace

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-20-d8d43f7fc947> in <module>
----> 1 movie_id_weights = tfra.dynamic_embedding.shadow_ops.embedding_lookup(movie_embeddings_shadow, ragged_movie_ids)

/data/miniconda3/envs/env-3.8.8/lib/python3.8/site-packages/tensorflow_recommenders_addons/dynamic_embedding/python/ops/shadow_embedding_ops.py in embedding_lookup(shadow, ids, partition_strategy, name, validate_indices)
    241       containing the values from the params tensor(s) for keys in ids.
    242   """
--> 243   ids = ops.convert_to_tensor(ids)
    244 
    245   if shadow.ids.dtype != ids.dtype:

/data/miniconda3/envs/env-3.8.8/lib/python3.8/site-packages/tensorflow/python/profiler/trace.py in wrapped(*args, **kwargs)
    181         with Trace(trace_name, **trace_kwargs):
    182           return func(*args, **kwargs)
--> 183       return func(*args, **kwargs)
    184 
    185     return wrapped

/data/miniconda3/envs/env-3.8.8/lib/python3.8/site-packages/tensorflow/python/framework/ops.py in convert_to_tensor(value, dtype, name, as_ref, preferred_dtype, dtype_hint, ctx, accepted_result_types)
   1693 
   1694     if ret is None:
-> 1695       ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
   1696 
   1697     if ret is NotImplemented:

/data/miniconda3/envs/env-3.8.8/lib/python3.8/site-packages/tensorflow/python/framework/constant_op.py in _constant_tensor_conversion_function(v, dtype, name, as_ref)
    341                                          as_ref=False):
    342   _ = as_ref
--> 343   return constant(v, dtype=dtype, name=name)
    344 
    345 

/data/miniconda3/envs/env-3.8.8/lib/python3.8/site-packages/tensorflow/python/framework/constant_op.py in constant(value, dtype, shape, name)
    265     ValueError: if called on a symbolic tensor.
    266   """
--> 267   return _constant_impl(value, dtype, shape, name, verify_shape=False,
    268                         allow_broadcast=True)
    269 

/data/miniconda3/envs/env-3.8.8/lib/python3.8/site-packages/tensorflow/python/framework/constant_op.py in _constant_impl(value, dtype, shape, name, verify_shape, allow_broadcast)
    277       with trace.Trace("tf.constant"):
    278         return _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
--> 279     return _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
    280 
    281   g = ops.get_default_graph()

/data/miniconda3/envs/env-3.8.8/lib/python3.8/site-packages/tensorflow/python/framework/constant_op.py in _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
    302 def _constant_eager_impl(ctx, value, dtype, shape, verify_shape):
    303   """Creates a constant on the current device."""
--> 304   t = convert_to_eager_tensor(value, ctx, dtype)
    305   if shape is None:
    306     return t

/data/miniconda3/envs/env-3.8.8/lib/python3.8/site-packages/tensorflow/python/framework/constant_op.py in convert_to_eager_tensor(value, ctx, dtype)
    100       dtype = dtypes.as_dtype(dtype).as_datatype_enum
    101   ctx.ensure_initialized()
--> 102   return ops.EagerTensor(value, ctx.device_name, dtype)
    103 
    104 

ValueError: TypeError: object of type 'RaggedTensor' has no len()

As a consequence, the high level keras embedding layers built on top of it failed as well, such as

• tfra.dynamic_embedding.keras.layers.BasicEmbedding
• tfra.dynamic_embedding.keras.layers.SquashedEmbedding

tf.RaggedTensor are usually used to represent features in recommender ranking models such as:

• the variable-length list of video_id viewed by a user in the past 3 days
• the variable-length list of author_id subscribed by a user in the past 24 hours

Is there any plan to support embedding lookup with ragged tensors?

Relevant information

• Are you willing to contribute it (yes/no): Yes
• Are you willing to maintain it going forward? (yes/no): Yes
• Is there a relevant academic paper? (if so, where): No
• Is there already an implementation in another framework? (if so, where): Not sure, haven't checked hugectr sok
• Was it part of tf.contrib? (if so, where): No

Which API type would this fall under (layer, metric, optimizer, etc.)

• tfra.dynamic_embedding.shadow_ops.embedding_lookup
• tfra.dynamic_embedding.keras.layers.BasicEmbedding
• tfra.dynamic_embedding.keras.layers.SquashedEmbedding

Who will benefit with this feature?
Any recommender developer who wishes to use ragged tensor to look up dynamic embedding

Any other info.

[Lifann]

@cockroachzl Thanks for feedback. Unfortunately, embedding_lookup currently only supports Tensor and SparseTensor. If using RaggedTensor, you'll have to process before lookup.

BTW, what do you expect on the lookup result, if using ragged tensor as input?

[cockroachzl]

@cockroachzl Thanks for feedback. Unfortunately, embedding_lookup currently only supports Tensor and SparseTensor. If using RaggedTensor, you'll have to process before lookup.

Got it, right now I work around this by using regular tf.tensor.

BTW, what do you expect on the lookup result, if using ragged tensor as input?

I expect the look up result to be the same as using tf.keras.layers.Embedding with ragged tensors as the input.

If the input is tf.ragged.constant([[1, 2], [3], [4, 5, 6]]), and the embedding size is 2, then the result of the embedding lookup is like

tf.ragged.constant([
    [[0.1,0.1], [0.2,0.2]], 
    [[0.3,0.3]], 
    [[0.4,0.4], [0.5,0.5], [0.6,0.6]]
])

[MoFHeka]

Maybe we need write an extra kernel code to lookup ragged tensor. If you could contribute it that will be so nice. @cockroachzl

[jq]

#397 @cockroachzl ragged tensor supported