def _process_batch(self, batch_data: List[tuple]) -> Tuple[Tensor, ...]:
"""
Preprocess the list of tuples representing a batched group of elements \
and return a tuple of arrays representing the batched values for every \
data attribute.
"""
unpacked = zip(*batch_data)
return tuple(judo.as_tensor(val) for val in unpacked)
def lennard_jones(x: np.ndarray) -> np.ndarray:
result = np.zeros(x.shape[0])
x = judo.to_numpy(x)
assert isinstance(x, np.ndarray)
for i in range(x.shape[0]):
try:
result[i] = _lennard_fast(x[i])
except ZeroDivisionError:
result[i] = np.inf
result = judo.as_tensor(result)
return result
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_iterable(self, x):
hashes = [self.hash_tensor(xi) for xi in x]
return judo.as_tensor(hashes, dtype=judo.dtype.hash_type)
def get_array_of_ids(self, n: int):
ids = numpy.arange(n) + self._seed + 1
self._seed += n + 1
return judo.as_tensor(ids)