Thanks,
yes the topic is more in general about what is the perspective when the compositional path doesn’t perform well.
Do we need to interact more strictly with the compiler team on the TF side before introducing a custom ops (often it is hard to collect a feedback)? I think new ops are interesting use cases to stress-test the compositional and the compiler stack transformations.
Will we have a new way to use the new compiler and runtime infra to write more portable high level custom ops?
If we are in a python only ecosystem repo, like keras*, where we need to contribute these “missing pieces”?
P.s.
For the embedding bag case (Addons, Pytorch TPU, JAX) at some point we had a sparse proposal  at:
  

      github.com/tensorflow/addons
  

  
    
  
    
  

  
    
      EmbeddingBag op and layer
    

    
      tensorflow:master ← Rocketknight1:master
    

    
      
        opened 
        
          
        
        Jan 19, 2021
      
      

      
        
          
          Rocketknight1
        
      

      
        
          +1100
          -0
        
      
    
  


  
    # Description

Brief Description of the PR:
This is a PR for the EmbeddingBag… op. Please don't merge it yet! Although it works, testing is incomplete and the file structure needs to be cleaned up. I'm opening it now just to get some initial feedback. I'll keep working on several of these issues (particularly 1, 3, 4 and 6 see below), but I'll need some feedback on 2) and 5), plus any other feedback you have for the rest of it!

Fixes # (issue)
#2201 

## Type of change

New layer and associated C++/CUDA op

# Comments

There are a few issues that need to be resolved before I'd feel comfortable with this being merged. In no particular order, they are:

1) The CUDA/C++ code is split with the forward and backward passes in separate files, which is not how other Tensorflow or Addons ops do it. This is just a style thing - I'll merge them soon. 

2) There are really two different entrypoints for users here, the function/op (analogous to tf.gather) and the layer (analogous to tf.keras.layers.Embedding). Like Embedding, the layer instantiates its own embeddings tensor and expects to be passed only indices and weights, whereas the function needs to be passed embeddings as well. Following PyTorch's naming conventions, I called the op embeddingbag and the layer EmbeddingBag, but this is almost certainly not what you want. What is the right way to name these two? Should I make the function/op a stateless Layer rather than just a function?

3) No support for float16/bfloat16 yet.

4) Because context->AllocateTemp continuously segfaulted for me when I was compiling in the custom-op repo, I used AllocateOutput to make some dummy outputs and then just used them as temp arrays. Compiling in tensorflow_addons itself seems much more stable, but I still need to go back and set that properly to AllocateTemp.

5) The CUDA/C++ ops expect a weight tensor. When no weights are passed, the Python wrapper instantiates dummy weights with `tf.ones_like()`. Is this acceptable?

6) More tests! I don't have any gradient tests at all yet, and I should probably add additional tests with weird shapes.
  

  

  
    
    
  

  


But then the custom ops was merged in Addons  (+1.100 lines for CPU/CUDA)