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)