torchrl 0.11.0__cp314-cp314-manylinux_2_28_aarch64.whl

torchrl/data/map/hash.py

@@ -0,0 +1,185 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+from torch.nn import Module
+
+
+class BinaryToDecimal(Module):
+    """A Module to convert binaries encoded tensors to decimals.
+
+    This is a utility class that allow to convert a binary encoding tensor (e.g. `1001`) to
+    its decimal value (e.g. `9`)
+
+    Args:
+        num_bits (int): the number of bits to use for the bases table.
+            The number of bits must be lower or equal to the input length and the input length
+            must be divisible by ``num_bits``. If ``num_bits`` is lower than the number of
+            bits in the input, the end result will be aggregated on the last dimension using
+            :func:`~torch.sum`.
+        device (torch.device): the device where inputs and outputs are to be expected.
+        dtype (torch.dtype): the output dtype.
+        convert_to_binary (bool, optional): if ``True``, the input to the ``forward``
+            method will be cast to a binary input using :func:`~torch.heavyside`.
+            Defaults to ``False``.
+
+    Examples:
+        >>> binary_to_decimal = BinaryToDecimal(
+        ...    num_bits=4, device="cpu", dtype=torch.int32, convert_to_binary=True
+        ... )
+        >>> binary = torch.Tensor([[0, 0, 1, 0, 0, 0, 0, 1], [0, 0, 0, 0, 0, 0, 10, 0]])
+        >>> decimal = binary_to_decimal(binary)
+        >>> assert decimal.shape == (2,)
+        >>> assert (decimal == torch.Tensor([3, 2])).all()
+    """
+
+    def __init__(
+        self,
+        num_bits: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        convert_to_binary: bool = False,
+    ):
+        super().__init__()
+        self.convert_to_binary = convert_to_binary
+        self.bases = 2 ** torch.arange(num_bits - 1, -1, -1, device=device, dtype=dtype)
+        self.num_bits = num_bits
+        self.zero_tensor = torch.zeros((1,), device=device)
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:
+        num_features = features.shape[-1]
+        if self.num_bits > num_features:
+            raise ValueError(f"{num_features=} is less than {self.num_bits=}")
+        elif num_features % self.num_bits != 0:
+            raise ValueError(f"{num_features=} is not divisible by {self.num_bits=}")
+
+        binary_features = (
+            torch.heaviside(features, self.zero_tensor)
+            if self.convert_to_binary
+            else features
+        )
+        feature_parts = binary_features.reshape(shape=(-1, self.num_bits))
+        digits = torch.vmap(torch.dot, (None, 0))(
+            self.bases, feature_parts.to(self.bases.dtype)
+        )
+        digits = digits.reshape(shape=(-1, features.shape[-1] // self.num_bits))
+        aggregated_digits = torch.sum(digits, dim=-1)
+        return aggregated_digits
+
+
+class SipHash(Module):
+    """A Module to Compute SipHash values for given tensors.
+
+    A hash function module based on SipHash implementation in python. Input tensors should have shape ``[batch_size, num_features]``
+    and the output shape will be ``[batch_size]``.
+
+    Args:
+        as_tensor (bool, optional): if ``True``, the bytes will be turned into integers
+            through the builtin ``hash`` function and mapped to a tensor. Default: ``True``.
+
+        .. warning:: This module relies on the builtin ``hash`` function.
+            To get reproducible results across runs, the ``PYTHONHASHSEED`` environment
+            variable must be set before the code is run (changing this value during code
+            execution is without effect).
+
+    Examples:
+        >>> # Assuming we set PYTHONHASHSEED=0 prior to running this code
+        >>> a = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]])
+        >>> b = a.clone()
+        >>> hash_module = SipHash(as_tensor=True)
+        >>> hash_a = hash_module(a)
+        >>> hash_a
+        tensor([-4669941682990263259, -3778166555168484291, -9122128731510687521])
+        >>> hash_b = hash_module(b)
+        >>> assert (hash_a == hash_b).all()
+    """
+
+    def __init__(self, as_tensor: bool = True):
+        super().__init__()
+        self.as_tensor = as_tensor
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor | list[bytes]:
+        hash_values = []
+        if x.dtype in (torch.bfloat16,):
+            x = x.to(torch.float16)
+        for x_i in x.detach().cpu().numpy():
+            hash_value = x_i.tobytes()
+            hash_values.append(hash_value)
+        if not self.as_tensor:
+            return hash_values
+        result = torch.tensor([hash(x) for x in hash_values], dtype=torch.int64)
+        return result
+
+
+class RandomProjectionHash(SipHash):
+    """A module that combines random projections with SipHash to get a low-dimensional tensor, easier to embed through :class:`~.SipHash`.
+
+    This module requires sklearn to be installed.
+
+    Keyword Args:
+        n_components (int, optional): the low-dimensional number of components of the projections.
+            Defaults to 16.
+        dtype_cast (torch.dtype, optional): the dtype to cast the projection to.
+            Defaults to ``torch.bfloat16``.
+        as_tensor (bool, optional): if ``True``, the bytes will be turned into integers
+            through the builtin ``hash`` function and mapped to a tensor. Default: ``True``.
+
+        .. warning:: This module relies on the builtin ``hash`` function.
+            To get reproducible results across runs, the ``PYTHONHASHSEED`` environment
+            variable must be set before the code is run (changing this value during code
+            execution is without effect).
+
+        init_method: TODO
+    """
+
+    _N_COMPONENTS_DEFAULT = 16
+
+    def __init__(
+        self,
+        *,
+        n_components: int | None = None,
+        dtype_cast=torch.bfloat16,
+        as_tensor: bool = True,
+        init_method: Callable[[torch.Tensor], torch.Tensor | None] | None = None,
+        **kwargs,
+    ):
+        if n_components is None:
+            n_components = self._N_COMPONENTS_DEFAULT
+
+        super().__init__(as_tensor=as_tensor)
+        self.register_buffer("_n_components", torch.as_tensor(n_components))
+
+        self._init = False
+        if init_method is None:
+            init_method = torch.nn.init.normal_
+        self.init_method = init_method
+
+        self.dtype_cast = dtype_cast
+        self.register_buffer("transform", torch.nn.UninitializedBuffer())
+
+    @property
+    def n_components(self):
+        return self._n_components.item()
+
+    def fit(self, x):
+        """Fits the random projection to the input data."""
+        self.transform.materialize(
+            (x.shape[-1], self.n_components), dtype=self.dtype_cast, device=x.device
+        )
+        self.init_method(self.transform)
+        self._init = True
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not self._init:
+            self.fit(x)
+        elif not self._init:
+            raise RuntimeError(
+                f"The {type(self).__name__} has not been initialized. Call fit before calling this method."
+            )
+        x = x.to(self.dtype_cast) @ self.transform
+        return super().forward(x)

torchrl/data/map/query.py

@@ -0,0 +1,204 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+
+from copy import deepcopy
+from typing import Any, TypeVar
+
+import torch
+import torch.nn as nn
+from tensordict import NestedKey, TensorDictBase
+from tensordict.nn.common import TensorDictModuleBase
+from torchrl._utils import logger as torchrl_logger
+from torchrl.data.map.hash import SipHash
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class HashToInt(nn.Module):
+    """Converts a hash value to an integer that can be used for indexing a contiguous storage."""
+
+    def __init__(self):
+        super().__init__()
+        self._index_to_index = {}
+
+    def __call__(self, key: torch.Tensor, extend: bool = False) -> torch.Tensor:
+        result = []
+        if extend:
+            for _item in key.tolist():
+                result.append(
+                    self._index_to_index.setdefault(_item, len(self._index_to_index))
+                )
+        else:
+            for _item in key.tolist():
+                result.append(
+                    self._index_to_index.get(_item, len(self._index_to_index))
+                )
+        return torch.tensor(result, device=key.device, dtype=key.dtype)
+
+    def state_dict(self) -> dict[str, torch.Tensor]:
+        values = torch.tensor(self._index_to_index.values())
+        keys = torch.tensor(self._index_to_index.keys())
+        return {"keys": keys, "values": values}
+
+    def load_state_dict(
+        self, state_dict: Mapping[str, Any], strict: bool = True, assign: bool = False
+    ):
+        keys = state_dict["keys"]
+        values = state_dict["values"]
+        self._index_to_index = {
+            key: val for key, val in zip(keys.tolist(), values.tolist())
+        }
+
+
+class QueryModule(TensorDictModuleBase):
+    """A Module to generate compatible indices for storage.
+
+    A module that queries a storage and return required index of that storage.
+    Currently, it only outputs integer indices (torch.int64).
+
+    Args:
+        in_keys (list of NestedKeys): keys of the input tensordict that
+            will be used to generate the hash value.
+        index_key (NestedKey): the output key where the index value will be written.
+            Defaults to ``"_index"``.
+
+    Keyword Args:
+        hash_key (NestedKey): the output key where the hash value will be written.
+            Defaults to ``"_hash"``.
+        hash_module (Callable[[Any], int] or a list of these, optional): a hash
+            module similar to :class:`~tensordict.nn.SipHash` (default).
+            If a list of callables is provided, its length must equate the number of in_keys.
+        hash_to_int (Callable[[int], int], optional): a stateful function that
+            maps a hash value to a non-negative integer corresponding to an index in a
+            storage. Defaults to :class:`~torchrl.data.map.HashToInt`.
+        aggregator (Callable[[int], int], optional): a hash function to group multiple hashes
+            together. This argument should only be passed when there is more than one ``in_keys``.
+            If a single ``hash_module`` is provided but no aggregator is passed, it will take
+            the value of the hash_module. If no ``hash_module`` or a list of ``hash_modules`` is
+            provided but no aggregator is passed, it will default to ``SipHash``.
+        clone (bool, optional): if ``True``, a shallow clone of the input TensorDict will be
+            returned. This can be used to retrieve the integer index within the storage,
+            corresponding to a given input tensordict. This can be overridden at runtime by
+            providing the ``clone`` argument to the forward method.
+            Defaults to ``False``.
+
+    Examples:
+        >>> query_module = QueryModule(
+        ...     in_keys=["key1", "key2"],
+        ...     index_key="index",
+        ...     hash_module=SipHash(),
+        ... )
+        >>> query = TensorDict(
+        ...     {
+        ...         "key1": torch.Tensor([[1], [1], [1], [2]]),
+        ...         "key2": torch.Tensor([[3], [3], [2], [3]]),
+        ...         "other": torch.randn(4),
+        ...     },
+        ...     batch_size=(4,),
+        ... )
+        >>> res = query_module(query)
+        >>> # The first two pairs of key1 and key2 match
+        >>> assert res["index"][0] == res["index"][1]
+        >>> # The last three pairs of key1 and key2 have at least one mismatching value
+        >>> assert res["index"][1] != res["index"][2]
+        >>> assert res["index"][2] != res["index"][3]
+
+    """
+
+    def __init__(
+        self,
+        in_keys: list[NestedKey],
+        index_key: NestedKey = "_index",
+        hash_key: NestedKey = "_hash",
+        *,
+        hash_module: Callable[[Any], int] | list[Callable[[Any], int]] | None = None,
+        hash_to_int: Callable[[int], int] | None = None,
+        aggregator: Callable[[Any], int] = None,
+        clone: bool = False,
+    ):
+        if len(in_keys) == 0:
+            raise ValueError("`in_keys` cannot be empty.")
+        in_keys = in_keys if isinstance(in_keys, list) else [in_keys]
+
+        super().__init__()
+        in_keys = self.in_keys = in_keys
+        self.out_keys = [index_key, hash_key]
+        index_key = self.out_keys[0]
+        self.hash_key = self.out_keys[1]
+
+        if aggregator is not None and len(self.in_keys) == 1:
+            torchrl_logger.warn(
+                "An aggregator was provided but there is only one in-key to be read. "
+                "This module will be ignored."
+            )
+        elif aggregator is None:
+            if hash_module is not None and not isinstance(hash_module, list):
+                aggregator = hash_module
+            else:
+                aggregator = SipHash()
+        if hash_module is None:
+            hash_module = [SipHash() for _ in range(len(self.in_keys))]
+        elif not isinstance(hash_module, list):
+            try:
+                hash_module = [
+                    deepcopy(hash_module) if len(self.in_keys) > 1 else hash_module
+                    for _ in range(len(self.in_keys))
+                ]
+            except Exception as err:
+                raise RuntimeError(
+                    "failed to deepcopy the hash module. Please provide a list of hash modules instead."
+                ) from err
+        elif len(hash_module) != len(self.in_keys):
+            raise ValueError(
+                "The number of hash_modules must match the number of in_keys. "
+                f"Got {len(hash_module)} hash modules but {len(in_keys)} in_keys."
+            )
+        if hash_to_int is None:
+            hash_to_int = HashToInt()
+
+        self.aggregator = aggregator
+        self.hash_module = dict(zip(self.in_keys, hash_module))
+        self.hash_to_int = hash_to_int
+
+        self.index_key = index_key
+        self.clone = clone
+
+    def forward(
+        self,
+        tensordict: TensorDictBase,
+        *,
+        extend: bool = True,
+        write_hash: bool = True,
+        clone: bool | None = None,
+    ) -> TensorDictBase:
+        hash_values = []
+
+        for k in self.in_keys:
+            hash_values.append(self.hash_module[k](tensordict.get(k)))
+        if len(self.in_keys) > 1:
+            hash_values = torch.stack(
+                hash_values,
+                dim=-1,
+            )
+            hash_values = self.aggregator(hash_values)
+        else:
+            hash_values = hash_values[0]
+
+        td_hash_value = self.hash_to_int(hash_values, extend=extend)
+
+        clone = clone if clone is not None else self.clone
+        if clone:
+            output = tensordict.copy()
+        else:
+            output = tensordict
+
+        output.set(self.index_key, td_hash_value)
+        if write_hash:
+            output.set(self.hash_key, hash_values)
+        return output