def _shuffle_layer(inputs, shuffle_fn):
  """Shuffles the elements according to bitwise left or right rotation.

  Args:
    inputs: Tensor input from previous layer
    shuffle_fn: Shift function rol or ror

  Returns:
    tf.Tensor: Inputs shifted according to shuffle_fn
  """
  seq_length = inputs.shape[1]
  n_bits = np.int32(np.log(seq_length - 1) / np.log(2.0)) + 1

  indices = np.arange(0, seq_length).astype('int32')
  rev_indices = shuffle_fn(indices, n_bits)
  return jnp.take(inputs, rev_indices, axis=1, mode='clip')
Shuffles the elements according to bitwise left or right rotation.

  Args:
    inputs: Tensor input from previous layer
    shuffle_fn: Shift function rol or ror

  Returns:
    tf.Tensor: Inputs shifted according to shuffle_fn
  
trax/models/research/rse.py