def hash_vecs(vecs, n_buckets_in, n_hashes, rng):
  """Hash vectors into buckets.

  Args:
    vecs: vectors to hash, a tensor of shape [batch_size, depth]
    n_buckets_in: an int or a list of ints, number of hash buckets;
      if it is a list, we do hierarchical hashing as specified by the list
    n_hashes: number of hashes
    rng: random generator to use for hashing

  Returns:
    A pair (buckets, n_buckets) where buckets is a tensor of shape
    [n_hashes, batch_size] of integers -- the hash bucket IDs, and
    n_buckets is an int, the total number of hash buckets, equal to
    the product of all items in n_buckets_in.
  """
  # See https://arxiv.org/pdf/1509.02897.pdf
  # We sample a different random rotation for each round of hashing to
  # decrease the probability of hash misses.
  if isinstance(n_buckets_in, int):
    assert n_buckets_in % 2 == 0
    rot_size = n_buckets_in
    n_buckets = n_buckets_in
  else:
    # Factorize the hash if n_buckets_in is a list or tuple
    rot_size, n_buckets = 0, 1
    for factor in n_buckets_in:
      assert factor % 2 == 0
      rot_size += factor
      n_buckets *= factor

  rotations_shape = (vecs.shape[-1], n_hashes, rot_size // 2)
  random_rotations = fastmath.random.normal(rng, rotations_shape).astype(
      np.float32)
  if fastmath.is_backend(fastmath.Backend.JAX):
    rotated_vecs = np.einsum('tf,fhb->htb', vecs, random_rotations)
  else:
    random_rotations = np.reshape(random_rotations,
                                  [-1, n_hashes * (rot_size // 2)])
    rotated_vecs = np.dot(vecs, random_rotations)
    rotated_vecs = np.reshape(rotated_vecs, [-1, n_hashes, rot_size//2])
    rotated_vecs = np.transpose(rotated_vecs, (1, 0, 2))

  if isinstance(n_buckets_in, int) or len(n_buckets_in) == 1:
    rotated_vecs = np.concatenate([rotated_vecs, -rotated_vecs], axis=-1)
    buckets = np.argmax(rotated_vecs, axis=-1).astype(np.int32)
  else:
    # Get the buckets for them and combine.
    buckets, cur_sum, cur_product = None, 0, 1
    for factor in n_buckets_in:
      rv = rotated_vecs[..., cur_sum:cur_sum + (factor // 2)]
      cur_sum += factor // 2
      rv = np.concatenate([rv, -rv], axis=-1)
      if buckets is None:
        buckets = np.argmax(rv, axis=-1).astype(np.int32)
      else:
        buckets += cur_product * np.argmax(rv, axis=-1).astype(np.int32)
      cur_product *= factor

  return buckets, n_buckets  # buckets is now (n_hashes, batch_size)
Hash vectors into buckets.

  Args:
    vecs: vectors to hash, a tensor of shape [batch_size, depth]
    n_buckets_in: an int or a list of ints, number of hash buckets;
      if it is a list, we do hierarchical hashing as specified by the list
    n_hashes: number of hashes
    rng: random generator to use for hashing

  Returns:
    A pair (buckets, n_buckets) where buckets is a tensor of shape
    [n_hashes, batch_size] of integers -- the hash bucket IDs, and
    n_buckets is an int, the total number of hash buckets, equal to
    the product of all items in n_buckets_in.
  
trax/layers/research/efficient_attention.py