Recently we had a refresh over a Deformable convloution WIP PR 51 in Addons.
I’ve cherry-picked this as an example as this requires us to maintain almost 3k lines of new code in the repository.
This is maintainer-ship overhead is also quite similar to what we have with other custom kernels PRs.
As Addons is one of the few Ecosystem repositories to support custom (c++) ops and the related CI infra it is quite normal that we have this kind of proposed PRs.
But as the codeownership of these components it is generally not so stable over time we would like to not merge, as possible, these custom ops PRs also to achieve a more broad hardware coverage.
What are the alternatives? How we could collaborate when a compositional implementation has huge performance gaps?
Often this kind of issues are shared across the “extend” ecosystem like  e.g. for the EmbeddingBag:

  

      github.com/pytorch/xla
  

  
    
  
	  
  

  
    
      lowering embeddingbag to XLA 6
    

    
      
        opened 
        
          
        
        Aug 5, 2020
      
      


      
        
          
          shz0116
        
      
    

    
        nostale
        op lowering
    
  


  
    
The embeddingbag operation has not been lowered to XLA. I saw aten:embeddingba…g from the profiling.
  

  

  
    
    
  

  



  

      github.com/tensorflow/addons
  

  
    
  
	  
  

  
    
      EmbeddingBag and Product-Key Memory Layers 1
    

    
      
        opened 
        
          
        
        Oct 14, 2020
      
      


      
        
          
          Rocketknight1
        
      
    

    
        Feature Request
        layers
    
  


  
    **Describe the feature and the current behavior/state.**
FAIR have a cool paper… where they introduce [Product-Key Memory Layers](https://arxiv.org/abs/1907.05242) - these are layers that can add a huge number of parameters (100M-1B) to a network with a very minimal compute overhead.

Unfortunately, implementing them efficiently depends on the [EmbeddingBag layer](https://pytorch.org/docs/stable/generated/torch.nn.EmbeddingBag.html) from Pytorch. This layer basically does a gather op followed by a weighted sum across the final dimension of the gather indices.

It is trivial to implement this op as a composition of two or three ops in Tensorflow, but doing so requires you to materialize the output of the gather, which in the case of Product-Key Memory layers is enormous, and usually blows out my GPU RAM. By combining these ops into a single efficient call, EmbeddingBag avoids ever materializing the extremely large pre-sum gather output. There's no efficient way to do the same in Tensorflow without a custom op.

I've already gotten a CUDA and (single-threaded) CPU implementation of EmbeddingBag working locally using the custom-op repo and associated docker image. I've verified correctness by comparing outputs and gradients to those from the manual composition of ops, and speed and memory usage are vastly improved. I could also contribute a TF implementation of the Product-Key Memory layer itself if desired.

**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):  https://arxiv.org/abs/1907.05242
- Is there already an implementation in another framework? (if so, where): Yes, EmbeddingBag is already a PyTorch layer 
- Was it part of tf.contrib? (if so, where):

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

**Who will benefit with this feature?**
People who want to squeeze loads of parameters into their model while maintaining fast throughput and aren't worried about overfitting. The paper used it for big autoregressive NLP Transformers, but I suspect you could deploy it in a lot of other places too.

**Any other info.**
I have only implemented the portions of EmbeddingBag necessary for Product-Key Memory layers.
  

  

  
    
    
  

  


EmbeddingBag op and layer by Rocketknight1 · Pull Request #2352 · tensorflow/addons · GitHub (1k lines)

  

      github.com/tensorflow/tensorflow
  

  
    
  
	  
  

  
    
      embedding_lookup cause ran out of memory  1
    

    
      
        opened 
        
          
        
        Oct 5, 2020
      
      


      
        
          
          shz0116
        
      
    

    
        TF 2.3
        comp:tpus
        comp:xla
        stat:awaiting tensorflower
        type:bug
    
  


  
    I am running the following code to test embedding_lookup.
```python

# comman…d:
# python3 -m pdb embtest.py --features=1000 --nnz=30 --batch=128
#
# error:
# *** tensorflow.python.framework.errors_impl.ResourceExhaustedError: 
#     Ran out of memory in memory space vmem. It should not be possible to run out of vmem - please file a bug against XLA.
#
import tensorflow as tf
import numpy as np
import sys
import os
import time

def measure(params, sp_ids, steps, thr):
  res = tf.nn.embedding_lookup([params[0:thr],params[thr:]], sp_ids, None, name="TEST1")
  print("Finished test")
  return res

if __name__ == "__main__":

  import sys 
  import argparse

  parser = argparse.ArgumentParser(
     description="Measure the performance of tensorflow embeddingbag using tf.nn.embedding" )
  parser.add_argument("--features", type=int, default=10)
  parser.add_argument("--em", type=int, default=2)
  parser.add_argument("--nnz", type=int, default=2)
  parser.add_argument("--batch", type=int, default=4)
  parser.add_argument("--steps", type=int, default=1)
  parser.add_argument("--warmups", type=int, default=0)

  args = parser.parse_args()

  features     = args.features
  em           = args.em
  nnz          = args.nnz
  batch        = args.batch
  steps        = args.steps
  warmups      = args.warmups

  sp_ids = np.random.randint(0, features, (batch * nnz,))
  res = tf.zeros([batch, em])

  resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu="grpc://"+os.environ["TPU_IP"])
  tf.config.experimental_connect_to_cluster(resolver)
  tf.tpu.experimental.initialize_tpu_system(resolver)
  print("   ")
  tpus = tf.config.list_logical_devices('TPU')
  print("There are {} tpu logical devices".format(len(tpus)))
  print(tpus[0])

  with tf.device('TPU:0'):
    params = tf.random.uniform([features, em])
    res = measure(params, sp_ids, tf.constant(steps), features//2)
 
  print(res)
```
But got the following error:

```bash
hongzhang@shan-tf1:~$ python  embtest.py --features=1000 --nnz=30 --batch=128
Eager execution :  True
2020-10-05 08:23:42.244623: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN)to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2020-10-05 08:23:42.250601: I tensorflow/core/platform/profile_utils/cpu_utils.cc:104] CPU Frequency: 2300000000 Hz
2020-10-05 08:23:42.251595: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x4c1dde0 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2020-10-05 08:23:42.251631: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Host, Default Version
2020-10-05 08:23:42.263068: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:301] Initialize GrpcChannelCache for job worker -> {0 -> 10.178.175.58:8470}
2020-10-05 08:23:42.263113: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:301] Initialize GrpcChannelCache for job localhost -> {0 -> localhost:38651}
2020-10-05 08:23:42.279709: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:301] Initialize GrpcChannelCache for job worker -> {0 -> 10.178.175.58:8470}
2020-10-05 08:23:42.279743: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:301] Initialize GrpcChannelCache for job localhost -> {0 -> localhost:38651}
2020-10-05 08:23:42.280176: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:405] Started server with target: grpc://localhost:38651
   
There are 8 tpu logical devices
LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:7', device_type='TPU')
Traceback (most recent call last):
  File "embtest.py", line 84, in <module>
    t1 = measure(params, sp_ids, tf.constant(steps), features//2)
  File "embtest.py", line 15, in measure
    res = tf.nn.embedding_lookup([params[0:thr],params[thr:]], sp_ids, None, name="TEST1")
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py", line 201, in wrapper
    return target(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/embedding_ops.py", line 394, in embedding_lookup_v2
    return embedding_lookup(params, ids, "div", name, max_norm=max_norm)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py", line 201, in wrapper
    return target(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/embedding_ops.py", line 328, in embedding_lookup
    transform_fn=None)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/embedding_ops.py", line 246, in _embedding_lookup_and_transform
    ret.set_shape(ids.get_shape().concatenate(element_shape_s))
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py", line 1206, in set_shape
    if not self.shape.is_compatible_with(shape):
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py", line 1167, in shape
    self._tensor_shape = tensor_shape.TensorShape(self._shape_tuple())
tensorflow.python.framework.errors_impl.ResourceExhaustedError: Ran out of memory in memory space vmem. It should not be possible to run out of vmem - please file a bug against XLA.

Largest program allocations in vmem:

  XLA label: register allocator spill slots
  Allocation type: scoped

  XLA label: %concatenate.724 = f32[3840,2]{0,1:T(2,128)} concatenate(f32[1,2]{0,1:T(2,128)}, f32[3,2]{0,1:T(2,128)}, f32[5,2]{0,1:T(2,128)}, f32[1,2]{0,1:T(2,128)}, ...(+2400)), dimensions={0}
  Allocation type: scoped

  XLA label: %concatenate.724 = f32[3840,2]{0,1:T(2,128)} concatenate(f32[1,2]{0,1:T(2,128)}, f32[3,2]{0,1:T(2,128)}, f32[5,2]{0,1:T(2,128)}, f32[1,2]{0,1:T(2,128)}, ...(+2400)), dimensions={0}
  Allocation type: scoped

  XLA label: %concatenate.724 = f32[3840,2]{0,1:T(2,128)} concatenate(f32[1,2]{0,1:T(2,128)}, f32[3,2]{0,1:T(2,128)}, f32[5,2]{0,1:T(2,128)}, f32[1,2]{0,1:T(2,128)}, ...(+2400)), dimensions={0}
  Allocation type: scoped

  XLA label: %concatenate.724 = f32[3840,2]{0,1:T(2,128)} concatenate(f32[1,2]{0,1:T(2,128)}, f32[3,2]{0,1:T(2,128)}, f32[5,2]{0,1:T(2,128)}, f32[1,2]{0,1:T(2,128)}, ...(+2400)), dimensions={0}
  Allocation type: scoped

2020-10-05 08:23:59.826142: W tensorflow/core/distributed_runtime/eager/remote_tensor_handle_data.cc:76] Unable to destroy remote tensor handles. If you are running a tf.function, it usually indicates some op in the graph gets an error: Ran out of memory in memory space vmem. It should not be possible to run out of vmem - please file a bug against XLA.

Largest program allocations in vmem:

  XLA label: register allocator spill slots
  Allocation type: scoped

  XLA label: %concatenate.724 = f32[3840,2]{0,1:T(2,128)} concatenate(f32[1,2]{0,1:T(2,128)}, f32[3,2]{0,1:T(2,128)}, f32[5,2]{0,1:T(2,128)}, f32[1,2]{0,1:T(2,128)}, ...(+2400)), dimensions={0}
  Allocation type: scoped

  XLA label: %concatenate.724 = f32[3840,2]{0,1:T(2,128)} concatenate(f32[1,2]{0,1:T(2,128)}, f32[3,2]{0,1:T(2,128)}, f32[5,2]{0,1:T(2,128)}, f32[1,2]{0,1:T(2,128)}, ...(+2400)), dimensions={0}
  Allocation type: scoped

  XLA label: %concatenate.724 = f32[3840,2]{0,1:T(2,128)} concatenate(f32[1,2]{0,1:T(2,128)}, f32[3,2]{0,1:T(2,128)}, f32[5,2]{0,1:T(2,128)}, f32[1,2]{0,1:T(2,128)}, ...(+2400)), dimensions={0}
  Allocation type: scoped

  XLA label: %concatenate.724 = f32[3840,2]{0,1:T(2,128)} concatenate(f32[1,2]{0,1:T(2,128)}, f32[3,2]{0,1:T(2,128)}, f32[5,2]{0,1:T(2,128)}, f32[1,2]{0,1:T(2,128)}, ...(+2400)), dimensions={0}
  Allocation type: scoped
```

**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04):
os: Linux
os kernel version: #1 SMP Debian 4.19.146-1 (2020-09-17)
os release version: 4.19.0-11-cloud-amd64
os platform: Linux-4.19.0-11-cloud-amd64-x86_64-with-debian-10.6
linux distribution: ('debian', '10.6', '')
linux os distribution: ('debian', '10.6', '')
mac version: ('', ('', '', ''), '')
uname: uname_result(system='Linux', node='shan-tf1', release='4.19.0-11-cloud-amd64', version='#1 SMP Debian 4.19.146-1 (2020-09-17)', machine='x86_64', processor='')
architecture: ('64bit', 'ELF')
machine: x86_64

- TensorFlow installed from (source or binary):
- TensorFlow version (use command below):

tf.version.VERSION = 2.3.0-dev20200620
tf.version.GIT_VERSION = v1.12.1-34769-gfd2d4cdb70
tf.version.COMPILER_VERSION = 7.3.1 20180303

- Python version:
- Bazel version (if compiling from source):
- GCC/Compiler version (if compiling from source):
- CUDA/cuDNN version:
- GPU model and memory:
  

  

  
    
    
  

  



  

      github.com/google/jax
  

  
    
  
	  
  

  
    
      np.take and np.einsum aren't fused properly 3
    

    
      
        opened 
        
          
        
        May 26, 2020
      
      


      
        
          
          AranKomat
        
      
    

    
        P2 (eventual)
        performance
        xla_issue
    
  


  
    I'm trying to translate [Product Key Memory ](https://arxiv.org/abs/1907.05242) …in PyTorch into JAX, and this requires the translation of [nn.EmbeddingBag](https://pytorch.org/docs/master/generated/torch.nn.EmbeddingBag.html) with per_sample_weights, as I haven't found any counterpart in JAX (but if you know, please let me know). For this, I wrote scatter_weighted_sum, the weighted sum version of scatter_add, in hope that it'll be efficient with fusing. However, jit didn't fuse np.take, reshape and np.einsum properly, which resulted in a huge intermediate object. Since #1979 concluded this sort of ops will be fused on GPU, I was wondering what is causing this problem. If by any chance this isn't supported on GPU, should this work on TPU? I'm using JAX ver. 0.1.67 with various Colab GPUs.   


````python
hidden_dim = 512
n_keys = 512
batch = 2 ** 15
knn = 32
heads = 4
key = random.PRNGKey(0)
values = random.normal(key, (n_keys ** 2, hidden_dim))         
indices = random.randint(key, (batch*heads, knn), 0, n_keys ** 2)     
weights = random.normal(key, (batch*heads, knn))     
@jit
def scatter_weighted_sum(inputs, indices, weights): 
    num_bags = weights.shape[-1]
    dim = inputs.shape[-1]
    indices = indices.reshape(-1)
    tmp = inputs.take(indices, axis=0).reshape(-1, num_bags, dim)
    return np.einsum('ind, in -> id', tmp, weights)

````
  

  

  
    
    
  

  


Thanks,
Stefano