def _input_with_indice_as_values(length, dim):
  """Retuns np.array of size (1, length, dim) where x[0, a, b] = a."""
  positions = []
  for i in range(length):
    positions.append([i] * dim)
  positions_input = np.array(positions)
  positions_input = np.expand_dims(positions_input, axis=0)
  return positions_input
Retuns np.array of size (1, length, dim) where x[0, a, b] = a.
trax/models/research/rse_test.py