(vecs, n_buckets_in, n_hashes, rng)
return buckets, n_buckets
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.
Hash vectors into buckets.
Hash vectors into buckets .
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
def hash_vecs ( vecs , n_buckets_in , n_hashes , rng ) : # 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