def new_backend_tensor(
x,
dtype=None,
device=None,
requires_grad: bool = None,
copy: bool = False,
pin_memory: bool = False,
):
kwargs = update_with_backend_values(requires_grad=requires_grad,
device=device,
copy=copy,
dtype=dtype)
if Backend.is_numpy():
return new_numpy_array(x,
dtype=kwargs.get("dtype"),
copy=kwargs.get("copy"))
elif Backend.is_torch():
return new_torch_tensor(
x,
dtype=kwargs.get("dtype"),
copy=kwargs.get("copy"),
device=kwargs.get("device"),
requires_grad=kwargs.get("requires_grad"),
pin_memory=pin_memory,
)
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 copy_torch(x: torch.Tensor, requires_grad, device=None) -> torch.Tensor:
grad = requires_grad if requires_grad is not None else Backend.requires_grad(
)
device = torch.device(
device) if device is not None else Backend.get_device()
new_tensor = x.clone()
if device is not None and new_tensor.device == device:
new_tensor = new_tensor.to(device)
if grad is not None and not grad:
new_tensor = new_tensor.detach()
return new_tensor
def to_torch(x,
requires_grad: bool = None,
device: str = None,
copy: bool = False):
use_grad = requires_grad if requires_grad is not None else Backend.requires_grad(
)
device = device if device is not None else Backend.get_device()
if isinstance(x, torch.Tensor):
return (copy_torch(x, device=device, requires_grad=use_grad)
if copy else _assign_device_and_grad(
x, device=device, requires_grad=use_grad))
return new_torch_tensor(x, device=device, copy=copy)
def __new_getattr(name):
if name in DATA_TYPE_NAMES:
return getattr(_data_types, name)()
elif name in AVAILABLE_FUNCTIONS:
return getattr(API, name)
try:
return __old_getattr(name)
except AttributeError as e:
if Backend.is_numpy():
val = getattr(numpy, name)
return __backend_wrap(val) if callable(val) else val
elif Backend.is_torch():
val = getattr(torch, name)
return __backend_wrap(val) if callable(val) else val
raise e
def random_nodes_generator(self, return_children: bool = False) -> NodeDataGenerator:
"""
Return a generator that yields the data of all the nodes sampling them at random.
Each value yielded corresponds to the data associated with one node and the \
edge connecting to one of its children. The edge data will be assigned \
to the parent node of the edge. This means the edge data of each node \
contains the data associated with the transition to a child.
Args:
return_children: If ``True`` the data corresponding to the child of \
each node in the path will also be returned.
Yields:
tuple of dictionaries containing the data for each node and edge \
sampled at random.
If ``return_children`` is ``False`` it will yield (node_data, \
edge_data).
If ``return_children`` is ``True`` it will yield (node_data, \
edge_data, next_node_data).
"""
with Backend.use_backend("numpy"):
permutaion = random_state.permutation(list(self.data.nodes))
for next_node in permutaion:
if next_node == self.root_id:
continue
node = self.get_parent(next_node)
yield self._one_node_tuples(node, next_node, return_children)
def copy(x, requires_grad: bool = None):
if x is None:
return
if not dtype.is_tensor(x):
x = JudoTensor(x)
funcs = {
"numpy": lambda x: x.copy(),
"torch": lambda x: copy_torch(x, requires_grad),
}
return Backend.execute(x, funcs)
def iterate_batches(self, batch_size: int, as_dict: bool = True):
with Backend.use_backend("numpy"):
indexes = random_state.permutation(range(len(self)))
for i in range(0, len(self), batch_size):
batch_ix = indexes[i:i + batch_size]
data = tuple(
[getattr(self, val)[batch_ix] for val in self.names])
if as_dict:
yield dict(zip(self.names, data))
else:
yield data
def iterate_values(self,
randomize: bool = False) -> Iterable[Tuple[Tensor]]:
"""
Return a generator that yields a tuple containing the data of each state \
stored in the memory.
"""
indexes = range(len(self))
if randomize:
with Backend.use_backend("numpy"):
indexes = random_state.permutation(indexes)
for i in indexes:
yield tuple([getattr(self, val)[i] for val in self.names])
def get_function(name):
if name in functions.fractalai.AVAILABLE_FUNCTIONS:
return getattr(functions.fractalai, name)
elif name in functions.batching.AVAILABLE_FUNCTIONS:
return getattr(functions.batching, name)
elif name in functions.images.AVAILABLE_FUNCTIONS:
return getattr(functions.images, name)
elif name in functions.notebook.AVAILABLE_FUNCTIONS:
return getattr(functions.notebook, name)
elif Backend.is_numpy():
backend = functions.numpy
else:
backend = functions.pytorch
return getattr(backend, name)
def to_backend(x: "Tensor",
requires_grad: bool = None,
device: str = None,
copy: bool = None) -> Tensor:
kwargs = update_with_backend_values(requires_grad=requires_grad,
device=device,
copy=copy)
if Backend.is_numpy():
return to_numpy(x, kwargs.get("copy"))
return to_torch(
x,
requires_grad=kwargs.get("requires_grad"),
device=kwargs.get("device"),
copy=kwargs.get("copy"),
)
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 float64():
funcs = {
"numpy": lambda x: numpy.float64,
"torch": lambda x: torch.float64
}
return Backend.execute(None, funcs)
def int32():
funcs = {"numpy": lambda x: numpy.int32, "torch": lambda x: torch.int32}
return Backend.execute(None, funcs)
def uint8():
funcs = {"numpy": lambda x: numpy.uint8, "torch": lambda x: torch.uint8}
return Backend.execute(None, funcs)
def bool():
funcs = {"numpy": lambda x: numpy.bool_, "torch": lambda x: torch.bool}
return Backend.execute(None, funcs)
def uses_true_hash(self) -> bool:
return Backend.use_true_hash()
def permutation(x):
idx = torch.randperm(x.shape[0])
sample = x[idx].to(Backend.get_device()).detach()
return to_backend(sample)
def update_with_backend_values(**kwargs):
user_kwargs = {k: v for k, v in kwargs.items() if v is not None}
backend_kwargs = Backend.get_backend_state()
backend_kwargs.update(user_kwargs)
return backend_kwargs
def true_hash_tensor(cls, x):
funcs = {
"numpy": cls.hash_numpy,
"torch": cls.hash_torch,
}
return Backend.execute(x, funcs)
def as_tensor(x, *args, **kwargs):
funcs = {
"numpy": lambda x: numpy.ascontiguousarray(x, *args, **kwargs),
"torch": lambda x: torch.as_tensor(x, *args, **kwargs),
}
return Backend.execute(x, funcs)
def astype(x, dtype):
funcs = {
"numpy": lambda x: x.astype(dtype),
"torch": lambda x: x.to(dtype),
}
return Backend.execute(x, funcs)
def __getattr__(cls, item):
funcs = {
"numpy": lambda x: getattr(cls._numpy_random_state, x),
"torch": lambda x: getattr(cls._torch_random_state, x),
}
return Backend.execute(item, funcs)
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