return_trainable currently is not implemented when using tf.function

[andrew-zzz]

System information Linux version 3.10.0 TensorFlow version 2.8.3,binary TensorFlow-Recommenders-Addons version 0.6 binary Python version:3.8 *Is GPU used? (yes/no):no

Describe the bug model.save() error that NotImplementedError: return_trainable currently is not implemented when using tf.function, i have set return_trainable=False both in de.embedding_lookup and de.embedding_lookup_sparse

Code to reproduce the issue

import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dense
from tensorflow_recommenders_addons import dynamic_embedding as de 
# ratings = tfds.load("movielens/100k-ratings", split="train")
import pandas as pd
import os

mutivalue_size = 1
max_mutivalue = 10
embedding_size = 8

def data_handle(x):
    data = tf.io.decode_csv(x,[
        tf.constant(0,dtype=tf.int64),
        [" "],
        tf.constant(0,dtype=tf.int64)
    ])
    return data[0],data[1],data[2]

class NCFModel(tf.keras.Model):

    def __init__(self):
        super(NCFModel, self).__init__()
        self.d0 = Dense(
            256,
            activation='relu',
            kernel_initializer=tf.keras.initializers.RandomNormal(0.0, 0.1),
            bias_initializer=tf.keras.initializers.RandomNormal(0.0, 0.1))
        self.d1 = Dense(
            64,
            activation='relu',
            kernel_initializer=tf.keras.initializers.RandomNormal(0.0, 0.1),
            bias_initializer=tf.keras.initializers.RandomNormal(0.0, 0.1))
        self.d2 = Dense(
            1,
            kernel_initializer=tf.keras.initializers.RandomNormal(0.0, 0.1),
            bias_initializer=tf.keras.initializers.RandomNormal(0.0, 0.1))

        self.id_embeddings = de.get_variable(
            name="id_dynamic_embeddings",
            dim=embedding_size,
            initializer=tf.keras.initializers.RandomNormal(-1, 1))

        self.loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)

    def input_handle(self,batch):
        mutivalue_str = tf.reshape(batch[1], [-1])

        mutivalue_split_str = tf.strings.split(mutivalue_str, "|").to_sparse()

        mutivalue_split_slice_split = tf.strings.split(mutivalue_split_str.values, ':').to_sparse()

        mutivalue_ids, mutivalue_weight = tf.sparse.split(sp_input=mutivalue_split_slice_split, num_split=2,
                                                          axis=1)

        mutivalue_ids_int64 = tf.strings.to_number(mutivalue_ids.values, out_type=tf.int64)
        mutivalue_ids_output = tf.SparseTensor(indices=mutivalue_split_str.indices,
                                               values=mutivalue_ids_int64,
                                               dense_shape=[mutivalue_size, max_mutivalue])
        mutivalue_weight_float = tf.strings.to_number(mutivalue_weight.values, out_type=tf.float32)
        mutivalue_weight_output = tf.SparseTensor(indices=mutivalue_split_str.indices,
                                                  values=mutivalue_weight_float,
                                                  dense_shape=[mutivalue_size, max_mutivalue])

        return batch[0],mutivalue_ids_output,mutivalue_weight_output,batch[2]

    @tf.function(input_signature=[
                                 (tf.TensorSpec([None], name='label', dtype=tf.int64),
                                  tf.TensorSpec([None], name='ids', dtype=tf.string),
                                  tf.TensorSpec([None], name='id', dtype=tf.int64),
                                  )])
    def call(self, batch):
        label,mutivalue_ids_output,mutivalue_weight_output,id = self.input_handle(batch)

        id_embedding = de.embedding_lookup(
            params=self.id_embeddings,
            ids=id,
            name="id_embedding",
            return_trainable=False
        )  

        multi_embedding = de.embedding_lookup_sparse(
            self.id_embeddings,
            mutivalue_ids_output,
            mutivalue_weight_output,
            partition_strategy=None,
            name='embedding_lookup_sparse',
            combiner='mean',
            max_norm=None,
            return_trainable=False
        )

        embeddings = tf.concat([multi_embedding, id_embedding], axis=1)
        dnn = self.d0(embeddings)
        dnn = self.d1(dnn)
        dnn = self.d2(dnn)
        out = tf.reshape(dnn, shape=[-1])
        loss = self.loss(label, out)
        return loss

model = NCFModel()
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
optimizer = de.DynamicEmbeddingOptimizer(optimizer)

model = NCFModel()
input = [tf.constant([1,1],dtype=tf.int64),tf.constant(['1:4|2:3','1:4|2:3'],dtype=tf.string),tf.constant([1,2],dtype=tf.int64)]
model(input)

Provide a reproducible test case that is the bare minimum necessary to generate the problem.

Other info / logs

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-2-19d833b99a13> in <module>
    109 model = NCFModel()
    110 input = [tf.constant([1,1],dtype=tf.int64),tf.constant(['1:4|2:3','1:4|2:3'],dtype=tf.string),tf.constant([1,2],dtype=tf.int64)]
--> 111 model(input)
    112 
    113 

/home/data/common/anaconda2/envs/tf283/lib/python3.8/site-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs)
     65     except Exception as e:  # pylint: disable=broad-except
     66       filtered_tb = _process_traceback_frames(e.__traceback__)
---> 67       raise e.with_traceback(filtered_tb) from None
     68     finally:
     69       del filtered_tb

/home/data/common/anaconda2/envs/tf283/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py in autograph_handler(*args, **kwargs)
   1145           except Exception as e:  # pylint:disable=broad-except
   1146             if hasattr(e, "ag_error_metadata"):
-> 1147               raise e.ag_error_metadata.to_exception(e)
   1148             else:
   1149               raise

NotImplementedError: Exception encountered when calling layer "ncf_model_3" (type NCFModel).

in user code:

    File "<ipython-input-1-19d833b99a13>", line 84, in call  *
        multi_embedding = de.embedding_lookup_sparse(
    File "/home/data/common/anaconda2/envs/tf283/lib/python3.8/site-packages/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_ops.py", line 779, in embedding_lookup_sparse  *
        embeddings, trainable_ = embedding_lookup(
    File "/home/data/common/anaconda2/envs/tf283/lib/python3.8/site-packages/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_ops.py", line 607, in embedding_lookup  *
        raise NotImplementedError(

    NotImplementedError: return_trainable currently is not implemented when using tf.function. Please use `Variable.trainable_store` or `Variable.get_trainable_by_name` to access the shadow trainable variable if call `embedding_lookup` series APIs inside tf.function scope.

Call arguments received:
  • batch=['tf.Tensor(shape=(2,), dtype=int64)', 'tf.Tensor(shape=(2,), dtype=string)', 'tf.Tensor(shape=(2,), dtype=int64)']

[andrew-zzz]

@Lifann @sunshinenum @rhdong please help

[rhdong]

Hi @andrew-zzz, I just checked the Keras demo and tutorial of dynamic embedding and I confirmed they work well, Here is the demo's log: tfra0.6.0_amazon-video-games-keras-eager-6_17_2023.log

As we mentioned in README: the conda is not recommended: https://github.com/tensorflow/recommenders-addons#compatibility-with-tensorflow

BTW, you shouldn't say that to anyone in Github. You're not welcome to us.

[andrew-zzz]

Hi @andrew-zzz, I just checked the Keras demo and tutorial of dynamic embedding and I confirmed they work well, Here is the demo's log: tfra0.6.0_amazon-video-games-keras-eager-6_17_2023.log

As we mentioned in README: the conda is not recommended: https://github.com/tensorflow/recommenders-addons#compatibility-with-tensorflow

BTW, you shouldn't say that to anyone in Github. You're not welcome to us.

I am sorry for that I am broken down that time, thank you for your reply that's a good project and I will still use it

[MoFHeka]

Has the bug been fixed? May I close this issue? @andrew-zzz

[MoFHeka]

This issue has been solved by following the demo instructions.

[alykhantejani]

I am running into this issue and I'm not sure what the resolution was from this issue?

[MoFHeka]

Hi @andrew-zzz, I just checked the Keras demo and tutorial of dynamic embedding and I confirmed they work well, Here is the demo's log: tfra0.6.0_amazon-video-games-keras-eager-6_17_2023.log

As we mentioned in README: the conda is not recommended: https://github.com/tensorflow/recommenders-addons#compatibility-with-tensorflow

BTW, you shouldn't say that to anyone in Github. You're not welcome to us.

@alykhantejani This issue could be caused by the wrong way of use

[alykhantejani]

ok let me open a new issue as Im running into this.

[MoFHeka]

@alykhantejani It is better to provide a minimum reproducible code.