code: def get_gates(self, x):
      
  
      # The balance loss don't propagate to the rest of the network
      x = tf.stop_gradient(x)
      # [length, depth] * [depth, nb_vectors * replicat]
      x = tf.matmul(x, self.t_vectors)
      # [length, nb_vector * replicat]
      x = tf.sign(x)  # Get on which side of the hyperplane the keys are.
  
      # x = tf.reshape(x, [-1, nb_replicat, nb_vector])
      # [length, replicat, nb_vector] * [nb_vector, 2^nb_vector - 1]
  
      x = tf.matmul(x, self.t_group, transpose_b=True) / self.nb_hyperplanes
      # We get a similarity score for each of the group between [-1, 1]
      # [length, (replicat,) 2^nb_vector - 1]
      # Do an argmax to get the most likely group for each replicat
      x = tf.argmax(x, axis=-1)
      # [length(, replicat)]
      # One-hot for compatibility with the sparse dispatcher
      x = tf.one_hot(x, self.nb_buckets)
      # TODO(epot): Use a loss to force an even distribution
      return x
docs: Return the bucket id of the given tensor.
  
      Args:
        x (tf.Tensor): float32 of shape [length, depth]
  
      Returns:
        tf.Tensor: One-hot vector int64 of shape [heads, length, nb_buckets]
          containing the id of the bucket