def hash_type():
funcs = {
"numpy":
lambda x: numpy.dtype("<U64")
if Backend.use_true_hash() else numpy.uint64,
"torch":
lambda x: torch.int64,
}
return Backend.execute(None, funcs)
def to_node_id(x):
if Backend.is_numpy():
return str(x) if Backend.use_true_hash() else int(x)
elif Backend.is_torch():
return int(x)
def uses_true_hash(self) -> bool:
return Backend.use_true_hash()
class Hasher:
_true_hash = bool(Backend.use_true_hash())
def __init__(self, seed: int = 0):
self._seed = seed
if self._true_hash:
self.pool = Pool()
@property
def uses_true_hash(self) -> bool:
return self._true_hash
@staticmethod
def hash_numpy(x: numpy.ndarray) -> int:
"""Return a value that uniquely identifies a numpy array."""
x = x.astype("|S576") if x.dtype == "O" else x
return xxhash.xxh64_hexdigest(x.tobytes())
@staticmethod
def hash_torch(x):
bytes = judo.to_numpy(x).tobytes()
return xxhash.xxh32_intdigest(bytes)
@classmethod
def true_hash_tensor(cls, x):
funcs = {
"numpy": cls.hash_numpy,
"torch": cls.hash_torch,
}
return Backend.execute(x, funcs)
def get_one_id(self):
self._seed += 1
return self._seed
def get_array_of_ids(self, n: int):
ids = numpy.arange(n) + self._seed + 1
self._seed += n + 1
return judo.as_tensor(ids)
def hash_tensor(self, x):
if self._true_hash:
return self.true_hash_tensor(x)
return 0
def hash_walkers(self, x):
if self._true_hash:
# hashes = self.pool.map(self.true_hash_tensor, x)
hashes = [self.true_hash_tensor(x_i) for x_i in x]
return judo.as_tensor(hashes)
return self.get_array_of_ids(x.shape[0])
def hash_state(self, state):
if self._true_hash:
_hash = hash(
tuple([
self.hash_tensor(x)
if k in state._tensor_names else hash(x)
for k, x in state.items()
]))
return _hash
return 0