import hashlib


def partition(example, num_buckets):
  """ Partition examples into multiple buckets """
  def determine_hash(input_bytes):
    m = hashlib.blake2b(input_bytes, digest_size=4)
    address = 0
    for i, b in enumerate(m.digest()):
      address += b * (2 ** (i * 8))
    return address

  address = determine_hash(example)
  bucket = address % num_buckets
  return bucket