Support keras style programming in eager mode and enable tf.function

[Lifann]

Background

Some of the design on Variable in current TensorFlow built on prerequisite of that variables are dense. And there is a gap between sparse Variable and dense Variable, since they have different shape mechanism, initializer implication, updating logic, etc., with agreement of RFC. Sparse Variable can be assume to be a superset to dense one, with more functionality and complex working mechanism.

Currently, in graph mode, the dynamic_embedding works fine, but in eager mode, we lack of either performance and not eazy to use. This PR aims to help the situation with supporting tf.function on dynamic_embedding and provide keras programming support.

In past versions of dynamic_embedding, we create a local variable to store the embedding lookup result temporarily, by calling embedding_lookup. If it was called inside scope of tf.function, the temporary variables will be translated to EagerTensor in the outputs:

@tf.function
def foo(devar, ids):
  lookup_result, trainable = de.embedding_lookup(devar, ids, return_trainable=True)
  return lookup_result, trainable

...

some_funcs_with_calling_foo()
tape.gradient(loss, [trainable])  # This will raise an error, since trainable is a EagerTensor, not a Variable.

Same issues are also existed when optimizers try to create slots.

Changes

• Add a new class inherited from TrainableWrapper to lift the local variable out of the tf.function scope.
• Support keras layers on dynamic_embedding

[MoFHeka]

embedding_lookup_unique is a better choice for following Keras Embedding layer?

[Lifann]

It does! In dynamic_embedding_ops.py, the xxx_embedding_lookup APIs are wrappers of embedding_lookup. So I'm trying to reuse the codes in these wrappers, by applying code transform on de.embedding_lookup. It will take some time, maybe just making it unique_embedding_lookup on current keras layer is realistic.

[Lifann]

The CI was failed caused by keras-team/keras#15586

[Lifann]

Keras initializer may not only come from tf.keras.initializer, when Keras save and load model, it could come from keras.initializer. Also when Keras save model, it could be a functools.partial(). So in file "tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_variable.py", it may be better like this in _convert_anything_to_init function:

try:
  from keras.initializers import initializers_v2 as ikinit2
except ImportError:
  ikinit2 = None
  pass  # for compatible with TF < 2.3.x

def _convert_anything_to_init(self, raw_init, dim):
  try:
    init = raw_init.func # init may be a functools.partial
  except:
    init = raw_init
  valid_list = [init_ops.Initializer, init_ops_v2.Initializer]
  if kinit2 is not None:
    valid_list.append(kinit2.Initializer)
  if ikinit2 is not None:
    valid_list.append(ikinit2.Initializer) # init may come from keras.initializer
  valid_list = tuple(valid_list)
  while callable(init):
    if isinstance(init, valid_list):
      self.initializer = init
      init = init(shape=[1])
    else:
      try:
        init = init(shape=[1])
      except:
        init = init()
  try:
    init = array_ops.reshape(init, [dim])
  except:
    init = array_ops.fill([dim], array_ops.reshape(init, [-1])[0])
  init = math_ops.cast(init, dtype=self.value_dtype)
  return init

Accept

[rhdong]

LGTM