title: Transfer python/numpy indexing to Tensorflow and improve performance
body: <p>In an earlier question <a href="https://stackoverflow.com/q/63774297/12859833">here</a>, I asked for advice on faster item assignment to an array. Since then, I have done some progress, eg I expanded the recommended version to take care of 3-D arrays, which is intended to resemble the batch size of the later training data for a Neural Network:</p>
  <pre class="lang-py prettyprint-override"><code>import numpy as np
  import time
  
  batch_dim = 2
  first_dim = 5
  second_dim = 7
  depth_dim = 10
  
  upper_count = 5000
  
  toy_dict = {k:np.random.random_sample(size = depth_dim) for k in range(upper_count)}
  a = np.array(list(toy_dict.values()))
  
  def create_input_3d(orig_arr):
    print(&quot;Input shape:&quot;, orig_arr.shape)
    goal_arr = np.full(shape=(batch_dim, orig_arr.shape[1], orig_arr.shape[2], depth_dim), fill_value=1234, dtype=float)
  
    print(&quot;Goal shape:&quot;, goal_arr.shape)
  
    idx = np.indices(orig_arr.shape)
    print(&quot;Idx shape&quot;, idx.shape)
    goal_arr[idx[0], idx[1], idx[2]] = a[orig_arr[idx[0], idx[1], idx[2]]]
  
    return goal_arr
  
  orig_arr_three_dim = np.random.randint(0, upper_count, size=(batch_dim, first_dim, second_dim))
  orig_arr_three_dim.shape # (2,5,7)
  
  reshaped = create_input_3d(orig_arr_three_dim)
  
  </code></pre>
  <p>I then decided on creating a custom layer to increase performance and do the conversion on the fly (decreases memory):</p>
  <pre class="lang-py prettyprint-override"><code>import tensorflow as tf
  from tensorflow import keras
  import numpy as np
  
  #custom layer
  class CustLayer(keras.layers.Layer):
      def __init__(self, info_matrix, first_dim, second_dim, info_dim, batch_size):
          super(CustLayer, self).__init__()
          self.w = tf.Variable(
              initial_value=info_matrix,
              trainable=False,
              dtype=tf.dtypes.float32
          )
          self.info_dim = info_dim
          self.first_dim = first_dim
          self.second_dim = second_dim
          self.batch_size = batch_size
  
      def call(self, orig_arr):
  
          goal_arr = tf.Variable(tf.zeros(shape=(self.batch_size, self.first_dim, self.second_dim, self.info_dim), dtype=float))
  
          #loop-approach (slower)
          for example in tf.range(self.batch_size):
            for row in tf.range(self.first_dim):
              for col in tf.range(self.second_dim):
                goal_arr[example,row,col].assign(self.w[orig_arr[example, row, col]])
          
          return goal_arr
  
  upper_count = 50
  info_length = 10
  batch_size = 4
  
  first_dim = 5
  second_dim = 7
  info_dim = 10
  
  info_dict = {k:np.random.random_sample(size = info_length) for k in range(upper_count)} #toy dict that stores information about
  info_matrix = np.array(list(info_dict.values()))
  
  
  linear_layer = CustLayer(info_matrix, first_dim=first_dim, second_dim=second_dim, info_dim=info_dim, batch_size=batch_size)
  
  test = []
  for i in range(batch_size):
    test.append(np.random.randint(1,upper_count, size=(first_dim,second_dim)))
  
  test = np.asarray(test)
  test.shape # (4, 5, 7)
  
  y= linear_layer(test)
  y.shape # TensorShape([4, 5, 7, 10])
  </code></pre>
  <p>Since advanced indexing (as in my first posted code) did not work, I returned to naive for-loops - which is too slow.</p>
  <p>What I am looking for is a way to use the advanced indexing as seen in the first code snippet, and reprogram it tf-compatible. This later enables me to use GPUs for my learning.</p>
  <p>In short: Input is of shape <code>(batch_size, first_dim, second_dim)</code>, return shape is <code>(batch_size, first_dim, second_dim, info_dim)</code>, getting rid of slow for-loops. Thanks in advance.</p>
  <p>Other answers I have looked in:
  <a href="https://stackoverflow.com/q/37061808/12859833">from 2016</a>, <a href="https://riptutorial.com/tensorflow/example/29038/numpy-like-indexing-using-tensors" rel="nofollow noreferrer">also old tf</a></p>