Source code for torchhd.embeddings

#
# MIT License
#
# Copyright (c) 2023 Mike Heddes, Igor Nunes, Pere Vergés, Denis Kleyko, and Danny Abraham
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
import math
from typing import Type, Union, Optional, Literal, Callable
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.nn.parameter import Parameter

import torchhd.functional as functional
from torchhd.tensors.base import VSATensor
from torchhd.tensors.map import MAPTensor
from torchhd.tensors.fhrr import FHRRTensor, type_conversion as fhrr_type_conversion
from torchhd.tensors.hrr import HRRTensor
from torchhd.types import VSAOptions


__all__ = [
    "Empty",
    "Identity",
    "Random",
    "Level",
    "Thermometer",
    "Circular",
    "Projection",
    "Sinusoid",
    "Density",
    "FractionalPower",
]


[docs] class Empty(nn.Embedding): """Embedding wrapper around :func:`~torchhd.empty`. Class inherits from `Embedding <https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html>`_ and supports the same keyword arguments. Args: num_embeddings (int): the number of hypervectors to generate. embedding_dim (int): the dimensionality of the hypervectors. vsa: (``VSAOptions``, optional): specifies the hypervector type to be instantiated. Default: ``"MAP"``. dtype (``torch.dtype``, optional): the desired data type of returned tensor. Default: if ``None`` uses default of VSATensor. device (``torch.device``, optional): the desired device of returned tensor. Default: if ``None``, uses the current device for the default tensor type (see torch.set_default_tensor_type()). ``device`` will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types. requires_grad (bool, optional): If autograd should record operations on the returned tensor. Default: ``False``. Examples:: >>> emb = embeddings.Empty(4, 6) >>> idx = torch.LongTensor([0, 1, 3]) >>> emb(idx) MAPTensor([[0., 0., 0., 0., 0., 0.], [0., 0., 0., 0., 0., 0.], [0., 0., 0., 0., 0., 0.]]) >>> emb = embeddings.Empty(4, 6, "FHRR") >>> idx = torch.LongTensor([0, 1, 3]) >>> emb(idx) FHRRTensor([[0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j], [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j], [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j]]) """ __constants__ = [ "num_embeddings", "embedding_dim", "vsa", "padding_idx", "max_norm", "norm_type", "scale_grad_by_freq", "sparse", ] vsa: VSAOptions def __init__( self, num_embeddings: int, embedding_dim: int, vsa: VSAOptions = "MAP", requires_grad: bool = False, padding_idx: Optional[int] = None, max_norm: Optional[float] = None, norm_type: float = 2.0, scale_grad_by_freq: bool = False, sparse: bool = False, device=None, dtype=None, **kwargs, ) -> None: factory_kwargs = {"device": device, "dtype": dtype} # Have to call Module init explicitly in order not to use the Embedding init nn.Module.__init__(self) self.num_embeddings = num_embeddings self.embedding_dim = embedding_dim self.vsa = vsa self.vsa_kwargs = kwargs if padding_idx is not None: if padding_idx > 0: assert ( padding_idx < self.num_embeddings ), "Padding_idx must be within num_embeddings" elif padding_idx < 0: assert ( padding_idx >= -self.num_embeddings ), "Padding_idx must be within num_embeddings" padding_idx = self.num_embeddings + padding_idx self.padding_idx = padding_idx self.max_norm = max_norm self.norm_type = norm_type self.scale_grad_by_freq = scale_grad_by_freq self.sparse = sparse embeddings = functional.empty( num_embeddings, embedding_dim, self.vsa, **factory_kwargs, **self.vsa_kwargs ) # Have to provide requires grad at the creation of the parameters to # prevent errors when instantiating a non-float embedding self.weight = Parameter(embeddings, requires_grad=requires_grad) # we don't need to set the padding to empty because it is already empty. def reset_parameters(self) -> None: factory_kwargs = {"device": self.weight.device, "dtype": self.weight.dtype} with torch.no_grad(): embeddings = functional.empty( self.num_embeddings, self.embedding_dim, self.vsa, **factory_kwargs, **self.vsa_kwargs, ) self.weight.copy_(embeddings)
[docs] def forward(self, input: Tensor) -> Tensor: vsa_tensor = functional.get_vsa_tensor_class(self.vsa) return super().forward(input).as_subclass(vsa_tensor)
[docs] class Identity(nn.Embedding): """Embedding wrapper around :func:`~torchhd.identity`. Class inherits from `Embedding <https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html>`_ and supports the same keyword arguments. Args: num_embeddings (int): the number of hypervectors to generate. embedding_dim (int): the dimensionality of the hypervectors. vsa: (``VSAOptions``, optional): specifies the hypervector type to be instantiated. Default: ``"MAP"``. dtype (``torch.dtype``, optional): the desired data type of returned tensor. Default: if ``None`` uses default of VSATensor. device (``torch.device``, optional): the desired device of returned tensor. Default: if ``None``, uses the current device for the default tensor type (see torch.set_default_tensor_type()). ``device`` will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types. requires_grad (bool, optional): If autograd should record operations on the returned tensor. Default: ``False``. Examples:: >>> emb = embeddings.Identity(4, 6) >>> idx = torch.LongTensor([0, 1, 3]) >>> emb(idx) MAPTensor([[1., 1., 1., 1., 1., 1.], [1., 1., 1., 1., 1., 1.], [1., 1., 1., 1., 1., 1.]]) >>> emb = embeddings.Identity(4, 6, "HRR") >>> idx = torch.LongTensor([0, 1, 3]) >>> emb(idx) HRRTensor([[1., 0., 0., 0., 0., 0.], [1., 0., 0., 0., 0., 0.], [1., 0., 0., 0., 0., 0.]]) """ __constants__ = [ "num_embeddings", "embedding_dim", "vsa", "padding_idx", "max_norm", "norm_type", "scale_grad_by_freq", "sparse", ] vsa: VSAOptions def __init__( self, num_embeddings: int, embedding_dim: int, vsa: VSAOptions = "MAP", requires_grad: bool = False, padding_idx: Optional[int] = None, max_norm: Optional[float] = None, norm_type: float = 2.0, scale_grad_by_freq: bool = False, sparse: bool = False, device=None, dtype=None, **kwargs, ) -> None: factory_kwargs = {"device": device, "dtype": dtype} # Have to call Module init explicitly in order not to use the Embedding init nn.Module.__init__(self) self.num_embeddings = num_embeddings self.embedding_dim = embedding_dim self.vsa = vsa self.vsa_kwargs = kwargs if padding_idx is not None: if padding_idx > 0: assert ( padding_idx < self.num_embeddings ), "Padding_idx must be within num_embeddings" elif padding_idx < 0: assert ( padding_idx >= -self.num_embeddings ), "Padding_idx must be within num_embeddings" padding_idx = self.num_embeddings + padding_idx self.padding_idx = padding_idx self.max_norm = max_norm self.norm_type = norm_type self.scale_grad_by_freq = scale_grad_by_freq self.sparse = sparse embeddings = functional.identity( num_embeddings, embedding_dim, self.vsa, **factory_kwargs, **self.vsa_kwargs ) # Have to provide requires grad at the creation of the parameters to # prevent errors when instantiating a non-float embedding self.weight = Parameter(embeddings, requires_grad=requires_grad) self._fill_padding_idx_with_empty() def reset_parameters(self) -> None: factory_kwargs = {"device": self.weight.device, "dtype": self.weight.dtype} with torch.no_grad(): embeddings = functional.identity( self.num_embeddings, self.embedding_dim, self.vsa, **factory_kwargs, **self.vsa_kwargs, ) self.weight.copy_(embeddings) self._fill_padding_idx_with_empty() def _fill_padding_idx_with_empty(self) -> None: factory_kwargs = {"device": self.weight.device, "dtype": self.weight.dtype} if self.padding_idx is not None: with torch.no_grad(): empty = functional.empty( 1, self.embedding_dim, self.vsa, **factory_kwargs, **self.vsa_kwargs ) self.weight[self.padding_idx].copy_(empty.squeeze(0))
[docs] def forward(self, input: Tensor) -> Tensor: vsa_tensor = functional.get_vsa_tensor_class(self.vsa) return super().forward(input).as_subclass(vsa_tensor)
[docs] class Random(nn.Embedding): """Embedding wrapper around :func:`~torchhd.random`. Class inherits from `Embedding <https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html>`_ and supports the same keyword arguments. Args: num_embeddings (int): the number of hypervectors to generate. embedding_dim (int): the dimensionality of the hypervectors. vsa: (``VSAOptions``, optional): specifies the hypervector type to be instantiated. Default: ``"MAP"``. dtype (``torch.dtype``, optional): the desired data type of returned tensor. Default: if ``None`` uses default of ``VSATensor``. device (``torch.device``, optional): the desired device of returned tensor. Default: if ``None``, uses the current device for the default tensor type (see torch.set_default_tensor_type()). ``device`` will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types. requires_grad (bool, optional): If autograd should record operations on the returned tensor. Default: ``False``. Examples:: >>> emb = embeddings.Random(4, 6) >>> idx = torch.LongTensor([0, 1, 3]) >>> emb(idx) MAPTensor([[-1., 1., -1., 1., -1., -1.], [ 1., -1., -1., -1., 1., -1.], [ 1., -1., 1., 1., 1., 1.]]) >>> emb = embeddings.Random(4, 6, "BSC") >>> idx = torch.LongTensor([0, 1, 3]) >>> emb(idx) BSCTensor([[ True, False, False, False, False, True], [False, True, True, True, False, True], [False, False, True, False, False, False]]) """ __constants__ = [ "num_embeddings", "embedding_dim", "vsa", "padding_idx", "max_norm", "norm_type", "scale_grad_by_freq", "sparse", ] vsa: VSAOptions def __init__( self, num_embeddings: int, embedding_dim: int, vsa: VSAOptions = "MAP", requires_grad: bool = False, padding_idx: Optional[int] = None, max_norm: Optional[float] = None, norm_type: float = 2.0, scale_grad_by_freq: bool = False, sparse: bool = False, device=None, dtype=None, **kwargs, ) -> None: factory_kwargs = {"device": device, "dtype": dtype} # Have to call Module init explicitly in order not to use the Embedding init nn.Module.__init__(self) self.num_embeddings = num_embeddings self.embedding_dim = embedding_dim self.vsa = vsa self.vsa_kwargs = kwargs if padding_idx is not None: if padding_idx > 0: assert ( padding_idx < self.num_embeddings ), "Padding_idx must be within num_embeddings" elif padding_idx < 0: assert ( padding_idx >= -self.num_embeddings ), "Padding_idx must be within num_embeddings" padding_idx = self.num_embeddings + padding_idx self.padding_idx = padding_idx self.max_norm = max_norm self.norm_type = norm_type self.scale_grad_by_freq = scale_grad_by_freq self.sparse = sparse embeddings = functional.random( num_embeddings, embedding_dim, self.vsa, **factory_kwargs, **self.vsa_kwargs ) # Have to provide requires grad at the creation of the parameters to # prevent errors when instantiating a non-float embedding self.weight = Parameter(embeddings, requires_grad=requires_grad) self._fill_padding_idx_with_empty() def reset_parameters(self) -> None: factory_kwargs = {"device": self.weight.device, "dtype": self.weight.dtype} with torch.no_grad(): embeddings = functional.random( self.num_embeddings, self.embedding_dim, self.vsa, **factory_kwargs, **self.vsa_kwargs, ) self.weight.copy_(embeddings) self._fill_padding_idx_with_empty() def _fill_padding_idx_with_empty(self) -> None: factory_kwargs = {"device": self.weight.device, "dtype": self.weight.dtype} if self.padding_idx is not None: with torch.no_grad(): empty = functional.empty( 1, self.embedding_dim, self.vsa, **factory_kwargs, **self.vsa_kwargs ) self.weight[self.padding_idx].copy_(empty.squeeze(0))
[docs] def forward(self, input: Tensor) -> Tensor: vsa_tensor = functional.get_vsa_tensor_class(self.vsa) return super().forward(input).as_subclass(vsa_tensor)
[docs] class Level(nn.Embedding): """Embedding wrapper around :func:`~torchhd.level`. Class inherits from `Embedding <https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html>`_ and supports the same keyword arguments. Args: num_embeddings (int): the number of hypervectors to generate. embedding_dim (int): the dimensionality of the hypervectors. vsa: (``VSAOptions``, optional): specifies the hypervector type to be instantiated. Default: ``"MAP"``. low (float, optional): The lower bound of the real number range that the levels represent. Default: ``0.0`` high (float, optional): The upper bound of the real number range that the levels represent. Default: ``1.0`` randomness (float, optional): r-value to interpolate between level-hypervectors at ``0.0`` and random-hypervectors at ``1.0``. Default: ``0.0``. dtype (``torch.dtype``, optional): the desired data type of returned tensor. Default: if ``None`` uses default of ``VSATensor``. device (``torch.device``, optional): the desired device of returned tensor. Default: if ``None``, uses the current device for the default tensor type (see torch.set_default_tensor_type()). ``device`` will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types. requires_grad (bool, optional): If autograd should record operations on the returned tensor. Default: ``False``. Values outside the interval between low and high are clipped to the closed bound. Examples:: >>> emb = embeddings.Level(4, 6) >>> x = torch.rand(4) >>> x tensor([0.6444, 0.9286, 0.9225, 0.3675]) >>> emb(x) MAPTensor([[ 1., 1., 1., -1., 1., 1.], [ 1., 1., -1., 1., 1., 1.], [ 1., 1., -1., 1., 1., 1.], [ 1., 1., 1., -1., -1., 1.]]) >>> emb = embeddings.Level(4, 6, "BSC") >>> x = torch.rand(4) >>> x tensor([0.1825, 0.1541, 0.4435, 0.1512]) >>> emb(x) BSCTensor([[False, True, False, False, False, False], [False, True, False, False, False, False], [False, True, False, False, False, False], [False, True, False, False, False, False]]) """ __constants__ = [ "num_embeddings", "embedding_dim", "vsa", "low", "high", "randomness", "padding_idx", "max_norm", "norm_type", "scale_grad_by_freq", "sparse", ] low: float high: float randomness: float vsa: VSAOptions def __init__( self, num_embeddings: int, embedding_dim: int, vsa: VSAOptions = "MAP", low: float = 0.0, high: float = 1.0, randomness: float = 0.0, requires_grad: bool = False, max_norm: Optional[float] = None, norm_type: float = 2.0, scale_grad_by_freq: bool = False, sparse: bool = False, device=None, dtype=None, **kwargs, ) -> None: factory_kwargs = {"device": device, "dtype": dtype} # Have to call Module init explicitly in order not to use the Embedding init nn.Module.__init__(self) self.num_embeddings = num_embeddings self.embedding_dim = embedding_dim self.vsa = vsa self.vsa_kwargs = kwargs self.low = low self.high = high self.randomness = randomness self.padding_idx = None self.max_norm = max_norm self.norm_type = norm_type self.scale_grad_by_freq = scale_grad_by_freq self.sparse = sparse embeddings = functional.level( num_embeddings, embedding_dim, self.vsa, randomness=randomness, **factory_kwargs, **self.vsa_kwargs, ) # Have to provide requires grad at the creation of the parameters to # prevent errors when instantiating a non-float embedding self.weight = Parameter(embeddings, requires_grad=requires_grad) def reset_parameters(self) -> None: factory_kwargs = {"device": self.weight.device, "dtype": self.weight.dtype} with torch.no_grad(): embeddings = functional.level( self.num_embeddings, self.embedding_dim, self.vsa, randomness=self.randomness, **factory_kwargs, **self.vsa_kwargs, ) self.weight.copy_(embeddings)
[docs] def forward(self, input: Tensor) -> Tensor: index = functional.value_to_index( input, self.low, self.high, self.num_embeddings ) index = index.clamp(min=0, max=self.num_embeddings - 1) vsa_tensor = functional.get_vsa_tensor_class(self.vsa) return super().forward(index).as_subclass(vsa_tensor)
[docs] class Thermometer(nn.Embedding): """Embedding wrapper around :func:`~torchhd.thermometer`. Class inherits from `Embedding <https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html>`_ and supports the same keyword arguments. Args: num_embeddings (int): the number of hypervectors to generate. embedding_dim (int): the dimensionality of the hypervectors. vsa: (``VSAOptions``, optional): specifies the hypervector type to be instantiated. Default: ``"MAP"``. low (float, optional): The lower bound of the real number range that the levels represent. Default: ``0.0`` high (float, optional): The upper bound of the real number range that the levels represent. Default: ``1.0`` dtype (``torch.dtype``, optional): the desired data type of returned tensor. Default: if ``None`` uses default of VSATensor. device (``torch.device``, optional): the desired device of returned tensor. Default: if ``None``, uses the current device for the default tensor type (see torch.set_default_tensor_type()). ``device`` will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types. requires_grad (bool, optional): If autograd should record operations on the returned tensor. Default: ``False``. Values outside the interval between low and high are clipped to the closed bound. Examples:: >>> emb = embeddings.Thermometer(4, 6) >>> x = torch.rand(4) >>> x tensor([0.5295, 0.0618, 0.0675, 0.1750]) >>> emb(x) MAPTensor([[ 1., 1., 1., 1., -1., -1.], [-1., -1., -1., -1., -1., -1.], [-1., -1., -1., -1., -1., -1.], [ 1., 1., -1., -1., -1., -1.]]) >>> emb = embeddings.Thermometer(4, 6, "FHRR") >>> x = torch.rand(4) >>> x tensor([0.2668, 0.7668, 0.8083, 0.6247]) >>> emb(x) FHRRTensor([[ 1.+0.j, 1.+0.j, -1.+0.j, -1.+0.j, -1.+0.j, -1.+0.j], [ 1.+0.j, 1.+0.j, 1.+0.j, 1.+0.j, -1.+0.j, -1.+0.j], [ 1.+0.j, 1.+0.j, 1.+0.j, 1.+0.j, -1.+0.j, -1.+0.j], [ 1.+0.j, 1.+0.j, 1.+0.j, 1.+0.j, -1.+0.j, -1.+0.j]]) """ __constants__ = [ "num_embeddings", "embedding_dim", "vsa", "low", "high", "padding_idx", "max_norm", "norm_type", "scale_grad_by_freq", "sparse", ] low: float high: float vsa: VSAOptions def __init__( self, num_embeddings: int, embedding_dim: int, vsa: VSAOptions = "MAP", low: float = 0.0, high: float = 1.0, requires_grad: bool = False, max_norm: Optional[float] = None, norm_type: float = 2.0, scale_grad_by_freq: bool = False, sparse: bool = False, device=None, dtype=None, **kwargs, ) -> None: factory_kwargs = {"device": device, "dtype": dtype} # Have to call Module init explicitly in order not to use the Embedding init nn.Module.__init__(self) self.num_embeddings = num_embeddings self.embedding_dim = embedding_dim self.vsa = vsa self.vsa_kwargs = kwargs self.low = low self.high = high self.padding_idx = None self.max_norm = max_norm self.norm_type = norm_type self.scale_grad_by_freq = scale_grad_by_freq self.sparse = sparse embeddings = functional.thermometer( num_embeddings, embedding_dim, self.vsa, **factory_kwargs, **self.vsa_kwargs ) # Have to provide requires grad at the creation of the parameters to # prevent errors when instantiating a non-float embedding self.weight = Parameter(embeddings, requires_grad=requires_grad) def reset_parameters(self) -> None: factory_kwargs = {"device": self.weight.device, "dtype": self.weight.dtype} with torch.no_grad(): embeddings = functional.thermometer( self.num_embeddings, self.embedding_dim, self.vsa, **factory_kwargs, **self.vsa_kwargs, ) self.weight.copy_(embeddings)
[docs] def forward(self, input: Tensor) -> Tensor: index = functional.value_to_index( input, self.low, self.high, self.num_embeddings ) index = index.clamp(min=0, max=self.num_embeddings - 1) vsa_tensor = functional.get_vsa_tensor_class(self.vsa) return super().forward(index).as_subclass(vsa_tensor)
[docs] class Circular(nn.Embedding): """Embedding wrapper around :func:`~torchhd.circular`. Class inherits from `Embedding <https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html>`_ and supports the same keyword arguments. Args: num_embeddings (int): the number of hypervectors to generate. embedding_dim (int): the dimensionality of the hypervectors. vsa: (``VSAOptions``, optional): specifies the hypervector type to be instantiated. Default: ``"MAP"``. phase (float, optional): The zero offset of the real number periodic interval that the circular levels represent. Default: ``0.0`` period (float, optional): The period of the real number periodic interval that the circular levels represent. Default: ``2 * pi`` randomness (float, optional): r-value to interpolate between circular-hypervectors at ``0.0`` and random-hypervectors at ``1.0``. Default: ``0.0``. dtype (``torch.dtype``, optional): the desired data type of returned tensor. Default: if ``None`` uses default of ``VSATensor``. device (``torch.device``, optional): the desired device of returned tensor. Default: if ``None``, uses the current device for the default tensor type (see torch.set_default_tensor_type()). ``device`` will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types. requires_grad (bool, optional): If autograd should record operations on the returned tensor. Default: ``False``. Examples:: >>> emb = embeddings.Circular(4, 6) >>> angle = torch.tensor([0.0, 3.141, 6.282, 9.423]) >>> emb(angle) MAPTensor([[-1., -1., -1., -1., -1., 1.], [-1., -1., 1., 1., -1., -1.], [-1., -1., -1., -1., -1., 1.], [-1., -1., 1., 1., -1., -1.]]) >>> emb = embeddings.Circular(4, 6, "BSC") >>> angle = torch.tensor([0.0, 3.141, 6.282, 9.423]) >>> emb(angle) BSCTensor([[False, True, False, False, True, True], [False, False, False, False, False, True], [False, True, False, False, True, True], [False, False, False, False, False, True]]) """ __constants__ = [ "num_embeddings", "embedding_dim", "vsa", "phase", "period", "randomness", "padding_idx", "max_norm", "norm_type", "scale_grad_by_freq", "sparse", ] phase: float period: float randomness: float vsa_tensor: Type[VSATensor] def __init__( self, num_embeddings: int, embedding_dim: int, vsa: VSAOptions = "MAP", phase: float = 0.0, period: float = 2 * math.pi, randomness: float = 0.0, requires_grad: bool = False, max_norm: Optional[float] = None, norm_type: float = 2.0, scale_grad_by_freq: bool = False, sparse: bool = False, device=None, dtype=None, **kwargs, ) -> None: factory_kwargs = {"device": device, "dtype": dtype} # Have to call Module init explicitly in order not to use the Embedding init nn.Module.__init__(self) self.num_embeddings = num_embeddings self.embedding_dim = embedding_dim self.vsa = vsa self.vsa_kwargs = kwargs self.phase = phase self.period = period self.randomness = randomness self.padding_idx = None self.max_norm = max_norm self.norm_type = norm_type self.scale_grad_by_freq = scale_grad_by_freq self.sparse = sparse embeddings = functional.circular( num_embeddings, embedding_dim, self.vsa, randomness=randomness, **factory_kwargs, **self.vsa_kwargs, ) # Have to provide requires grad at the creation of the parameters to # prevent errors when instantiating a non-float embedding self.weight = Parameter(embeddings, requires_grad=requires_grad) def reset_parameters(self) -> None: factory_kwargs = {"device": self.weight.device, "dtype": self.weight.dtype} with torch.no_grad(): embeddings = functional.circular( self.num_embeddings, self.embedding_dim, self.vsa, randomness=self.randomness, **factory_kwargs, **self.vsa_kwargs, ) self.weight.copy_(embeddings)
[docs] def forward(self, input: Tensor) -> Tensor: mapped = functional.map_range( input, self.phase, self.period, 0, self.num_embeddings ) index = mapped.round().long() % self.num_embeddings vsa_tensor = functional.get_vsa_tensor_class(self.vsa) return super().forward(index).as_subclass(vsa_tensor)
[docs] class Projection(nn.Module): r"""Embedding using a random projection matrix. Implemented based on `A Theoretical Perspective on Hyperdimensional Computing <https://arxiv.org/abs/2010.07426>`_. It computes :math:`x \Phi^{\mathsf{T}}` where :math:`\Phi \in \mathbb{R}^{d \times m}` is a matrix whose rows are uniformly sampled at random from the surface of an :math:`d`-dimensional unit sphere. This encoding ensures that similarities in the input space are preserved in the hyperspace. Args: in_features (int): the dimensionality of the input feature vector. out_features (int): the dimensionality of the hypervectors. vsa: (``VSAOptions``, optional): specifies the hypervector type to be instantiated. Default: ``"MAP"``. requires_grad (bool, optional): If autograd should record operations on the returned tensor. Default: ``False``. dtype (``torch.dtype``, optional): the desired data type of returned tensor. Default: if ``None``, uses a global default (see ``torch.set_default_tensor_type()``). device (``torch.device``, optional): the desired device of returned tensor. Default: if ``None``, uses the current device for the default tensor type (see torch.set_default_tensor_type()). ``device`` will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types. Examples:: >>> embed = embeddings.Projection(6, 5) >>> x = torch.randn(3, 6) >>> x tensor([[ 0.4119, -0.4284, 1.8022, 0.3715, -1.4563, -0.2842], [-0.3772, -1.2664, -1.5173, 1.3317, 0.4707, -1.3362], [-1.8142, 0.0274, -1.0989, 0.8193, 0.7619, 0.9181]]) >>> embed(x).sign() MAPTensor([[-1., 1., 1., 1., 1.], [ 1., 1., 1., 1., 1.], [ 1., -1., -1., -1., -1.]]) """ __constants__ = ["in_features", "out_features", "vsa"] in_features: int out_features: int vsa: VSAOptions weight: torch.Tensor def __init__( self, in_features, out_features, vsa: VSAOptions = "MAP", requires_grad=False, device=None, dtype=None, ): factory_kwargs = {"device": device, "dtype": dtype} super(Projection, self).__init__() self.in_features = in_features self.out_features = out_features self.vsa = vsa if vsa not in {"MAP", "HRR", "VTB"}: raise ValueError( f"Projection embedding supports MAP, HRR, VTB but provided: {vsa}" ) self.weight = nn.parameter.Parameter( torch.empty((out_features, in_features), **factory_kwargs), requires_grad=requires_grad, ) self.reset_parameters() def reset_parameters(self) -> None: nn.init.normal_(self.weight, 0, 1) self.weight.data.copy_(F.normalize(self.weight.data))
[docs] def forward(self, input: torch.Tensor) -> torch.Tensor: vsa_tensor = functional.get_vsa_tensor_class(self.vsa) return F.linear(input, self.weight).as_subclass(vsa_tensor)
[docs] class Sinusoid(nn.Module): r"""Embedding using a nonlinear random projection Implemented based on `Scalable Edge-Based Hyperdimensional Learning System with Brain-Like Neural Adaptation <https://dl.acm.org/doi/abs/10.1145/3458817.3480958>`_. It computes :math:`\cos(x \Phi^{\mathsf{T}} + b) \odot \sin(x \Phi^{\mathsf{T}})` where :math:`\Phi \in \mathbb{R}^{d \times m}` is a matrix whose rows are uniformly sampled at random from the surface of an :math:`d`-dimensional unit sphere and :math:`b \in \mathbb{R}^{d}` is a vectors whose elements are sampled uniformly at random between 0 and :math:`2\pi`. Args: in_features (int): the dimensionality of the input feature vector. out_features (int): the dimensionality of the hypervectors. vsa: (``VSAOptions``, optional): specifies the hypervector type to be instantiated. Default: ``"MAP"``. requires_grad (bool, optional): If autograd should record operations on the returned tensor. Default: ``False``. dtype (``torch.dtype``, optional): the desired data type of returned tensor. Default: if ``None``, uses a global default (see ``torch.set_default_tensor_type()``). device (``torch.device``, optional): the desired device of returned tensor. Default: if ``None``, uses the current device for the default tensor type (see torch.set_default_tensor_type()). ``device`` will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types. Examples:: >>> embed = embeddings.Sinusoid(6, 5) >>> x = torch.randn(3, 6) >>> x tensor([[ 0.5043, 0.3161, -0.0938, 0.6134, -0.1280, 0.3647], [-0.1907, 1.6468, -0.3242, 0.8614, 0.3332, -0.2055], [-0.8662, -1.3861, -0.1577, 0.1321, -0.1157, -2.8928]]) >>> embed(x) MAPTensor([[-0.0555, 0.2292, -0.1833, 0.0301, -0.2416], [-0.0725, 0.7042, -0.5644, 0.2235, 0.3603], [-0.9021, 0.8899, -0.9802, 0.3565, 0.2367]]) """ __constants__ = ["in_features", "out_features", "vsa"] in_features: int out_features: int vsa: VSAOptions weight: torch.Tensor bias: torch.Tensor def __init__( self, in_features, out_features, vsa: VSAOptions = "MAP", requires_grad=False, device=None, dtype=None, ): factory_kwargs = {"device": device, "dtype": dtype} super(Sinusoid, self).__init__() self.in_features = in_features self.out_features = out_features self.vsa = vsa if vsa not in {"MAP", "HRR", "VTB"}: raise ValueError( f"Sinusoid embedding supports MAP, HRR, VTB but provided: {vsa}" ) self.weight = nn.parameter.Parameter( torch.empty((out_features, in_features), **factory_kwargs), requires_grad=requires_grad, ) self.bias = nn.parameter.Parameter( torch.empty((1, out_features), **factory_kwargs), requires_grad=requires_grad, ) self.reset_parameters() def reset_parameters(self) -> None: nn.init.normal_(self.weight, 0, 1) self.weight.data.copy_(F.normalize(self.weight.data)) nn.init.uniform_(self.bias, 0, 2 * math.pi)
[docs] def forward(self, input: torch.Tensor) -> torch.Tensor: projected = F.linear(input, self.weight) output = torch.cos(projected + self.bias) * torch.sin(projected) vsa_tensor = functional.get_vsa_tensor_class(self.vsa) return output.as_subclass(vsa_tensor)
[docs] class Density(nn.Module): """Performs the transformation of input data into hypervectors according to the intRVFL model. See details in `Density Encoding Enables Resource-Efficient Randomly Connected Neural Networks <https://doi.org/10.1109/TNNLS.2020.3015971>`_. Args: in_features (int): the dimensionality of the input feature vector. out_features (int): the dimensionality of the hypervectors. vsa: (``VSAOptions``, optional): specifies the hypervector type to be instantiated. Default: ``"MAP"``. low (float, optional): The lower bound of the real number range that the levels of the thermometer encoding represent. Default: ``0.0`` high (float, optional): The upper bound of the real number range that the levels of the thermometer encoding represent. Default: ``1.0`` dtype (``torch.dtype``, optional): the desired data type of returned tensor. Default: if ``None`` uses default of ``VSATensor``. device (``torch.device``, optional): the desired device of returned tensor. Default: if ``None``, uses the current device for the default tensor type (see torch.set_default_tensor_type()). ``device`` will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types. requires_grad (bool, optional): If autograd should record operations on the returned tensor. Default: ``False``. Examples:: >>> embed = embeddings.Density(6, 5) >>> x = torch.randn(3, 6) >>> x tensor([[ 0.5430, 1.0740, 0.7250, -0.3410, -0.1318, 1.3188], [ 0.4373, 1.2400, -0.2264, 1.2448, -0.2040, -0.7831], [ 1.7460, -0.7359, -1.3271, 0.4338, -0.2401, 1.6553]]) >>> embed(x) MAPTensor([[ 2., 2., -2., -2., 0.], [ 4., 0., 6., 4., 0.], [ 4., -4., -2., -4., -4.]]) """ def __init__( self, in_features: int, out_features: int, vsa: VSAOptions = "MAP", low: float = 0.0, high: float = 1.0, device=None, dtype=None, requires_grad: bool = False, **kwargs, ): factory_kwargs = { "device": device, "dtype": dtype, "requires_grad": requires_grad, } super(Density, self).__init__() # A set of random vectors used as unique IDs for features of the dataset. self.key = Random(in_features, out_features, vsa, **factory_kwargs, **kwargs) # Thermometer encoding used for transforming input data. self.density_encoding = Thermometer( out_features + 1, out_features, vsa, low=low, high=high, **factory_kwargs, **kwargs, ) def reset_parameters(self) -> None: self.key.reset_parameters() self.density_encoding.reset_parameters() # Specify the steps needed to perform the encoding
[docs] def forward(self, input: Tensor) -> Tensor: # Perform binding of key and value vectors output = functional.bind(self.key.weight, self.density_encoding(input)) # Perform the superposition operation on the bound key-value pairs return functional.multibundle(output)
[docs] class FractionalPower(nn.Module): """Class for fractional power encoding (FPE) method that forms hypervectors for given values, kernel shape, bandwidth, and dimensionality. Implements similarity-preserving hypervectors approximating desired kernel shape as described in `Computing on Functions Using Randomized Vector Representations <https://arxiv.org/abs/2109.03429>`_. Args: in_features (int): the dimensionality of the input feature vector. out_features (int): the dimensionality of the hypervectors. distribution (str, optional): hyperparameter defining the shape of the kernel by specifying a particular probability distribution that is used to sample the base hypervector(s). Default: ``"sinc"``. bandwidth (float, optional): positive hyperparameter defining the width of the similarity kernel. Lower values lead to broader kernels while larger values lead to more narrow kernels. Default: ``1.0``. vsa: (``VSAOptions``, optional): specifies the hypervector type to be instantiated. Default: ``"FHRR"``. dtype (``torch.dtype``, optional): the desired data type of returned tensor. Default: if ``None`` depends on VSATensor. device (``torch.device``, optional): the desired device of returned tensor. Default: if ``None``, uses the current device for the default tensor type (see torch.set_default_tensor_type()). ``device`` will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types. requires_grad (bool, optional): If autograd should record operations on the returned tensor. Default: ``False``. Examples:: >>> embed = embeddings.FractionalPower(1, 6, "sinc", 1.0, "FHRR") >>> embed(torch.arange(1, 4, 1.).view(-1, 1)) FHRRTensor([[-0.7181-0.6959j, -0.5269+0.8499j, -0.0848+0.9964j, 0.9720-0.2348j, 0.6358+0.7718j, 0.4352+0.9003j], [ 0.0314+0.9995j, -0.4447-0.8957j, -0.9856-0.1689j, 0.8897-0.4565j, -0.1915+0.9815j, -0.6212+0.7836j], [ 0.6730-0.7396j, 0.9956+0.0940j, 0.2519-0.9678j, 0.7576-0.6527j, -0.8793+0.4762j, -0.9759-0.2183j]]) """ # The collection of distributions for basic predefined kernels predefined_kernels = { "sinc": torch.distributions.Uniform(-math.pi, math.pi), "gaussian": torch.distributions.Normal(0.0, 1.0), } def __init__( self, in_features: int, out_features: int, distribution: Union[ torch.distributions.Distribution, Literal["sinc", "gaussian"] ] = "sinc", bandwidth: float = 1.0, vsa: Literal["HRR", "FHRR"] = "FHRR", device=None, dtype=None, requires_grad: bool = False, ) -> None: super(FractionalPower, self).__init__() self.in_features = in_features # data dimensions self.out_features = out_features # hypervector dimensions self.bandwidth = bandwidth self.requires_grad = requires_grad if vsa not in {"HRR", "FHRR"}: raise ValueError( f"FractionalPower embedding only supports HRR and FHRR but provided: {vsa}" ) self.vsa_tensor = functional.get_vsa_tensor_class(vsa) # If a specific dtype is specified make sure it is supported by the VSA model if dtype != None and dtype not in self.vsa_tensor.supported_dtypes: raise ValueError(f"dtype {dtype} not supported by {vsa}") # The internal weights/phases are stored as floats even if the output is a complex tensor if dtype != None and vsa == "FHRR": dtype = fhrr_type_conversion[dtype] factory_kwargs = {"device": device, "dtype": dtype} # If the distribution is a string use the presets in predefined_kernels if isinstance(distribution, str): try: self.distribution = self.predefined_kernels[distribution] except KeyError: available_names = ", ".join(list(self.predefined_kernels.keys())) raise ValueError( f"{distribution} kernel is not supported, use one of: {available_names}, or provide a custom torch distribution." ) else: self.distribution = distribution # Initialize encoding's parameters self.weight = nn.Parameter( torch.empty(self.out_features, self.in_features, **factory_kwargs), requires_grad, ) self.reset_parameters() # Sample the angles using the provided distribution
[docs] def reset_parameters(self) -> None: """Generate the angles for basis hypervector(s) to be used for encoding the data.""" sample_shape = self.distribution.sample().shape # Check HD/VSA model type if self.vsa_tensor == FHRRTensor: # Generate the angles for base hypervector(s) that determines the shape of the FPE kernel # If the distribution is one-dimensional this implies that base hypervectors are independent so it is safe to generate self.in_features * self.out_features independent samples if sample_shape == (): # Draw angles from a uniform distribution for base hypervector(s). Note that data dimensions here are independent but this does not have to be always the case phases = self.distribution.sample((self.out_features, self.in_features)) phases = phases.to(self.weight) self.weight.data.copy_(phases) # If base hypervectors are correlated then the dimensionality of the distribution should match that of the data elif sample_shape == (self.in_features,): phases = self.distribution.sample((self.out_features,)) phases = phases.to(self.weight) self.weight.data.copy_(phases) # Raise error due to the ambiguity of the situation else: raise ValueError( f"The provided distribution has shape {sample_shape} while the input data expects shape () or ({self.in_features},) so there is a mismatch." ) elif self.vsa_tensor == HRRTensor: # Fewer angles are needed dimensions_real = int((self.out_features - 1) / 2) # Generate the angles for base hypervector(s) that determines the shape of the FPE kernel # If the distribution is one-dimensional this implies that base hypervectors are independent so it is safe to generate self.in_features * self.out_features independent samples if sample_shape == (): # Draw angles from a uniform distribution for base hypervector(s). Note that data dimensions here are independent but this does not have to be always the case phases = self.distribution.sample((dimensions_real, self.in_features)) # If base hypervectors are correlated then the dimensionality of the distribution should match that of the data elif sample_shape == (self.in_features,): phases = self.distribution.sample((dimensions_real,)) # Raise error due to the ambiguity of the situation else: raise ValueError( f"The provided distribution has shape {sample_shape} while the input data expects shape () or ({self.in_features},) so there is a mismatch." ) # Make the generated angles negatively symmetric so they look as a spectrum phases = torch.cat( ( phases, torch.zeros(1, self.in_features), -torch.flip(phases, dims=[0]), ), dim=0, ) if self.out_features % 2 == 0: phases = torch.cat((torch.zeros(1, self.in_features), phases), dim=0) phases = phases.to(self.weight) # Set the generated angles to the object's parameters self.weight.data.copy_(phases)
[docs] def basis(self): """Return the values of the base hypervector(s)""" # Use the angles in self.weight to obtain the values of the base hypervector(s) if self.vsa_tensor == FHRRTensor: hvs = torch.complex(self.weight.cos(), self.weight.sin()).T hvs = hvs.as_subclass(FHRRTensor) elif self.vsa_tensor == HRRTensor: complex_hv = torch.complex(self.weight.cos(), self.weight.sin()).T hvs = torch.real( torch.fft.ifft(torch.fft.ifftshift(complex_hv, dim=1), dim=1) ) hvs = hvs.as_subclass(HRRTensor) return hvs
[docs] def forward(self, input: Tensor) -> Tensor: """Creates a fractional power encoding (FPE) for given values. Args: input (Tensor): values for which FPE hypervectors should be generated. Either a vector or a batch of vectors. Shapes: - Input: :math:`(*, f)` where f is the in_features and * is an optional batch dimension. - Output: :math:`(*, d)` where d is the out_features and * is an optional batch dimension. """ # Perform FPE of the desired values using the base hypervector(s) # Simultaneously computes angles for given values and their sum that is equivalent to the binding if self.vsa_tensor == FHRRTensor: phases = F.linear(self.bandwidth * input, self.weight) hv = torch.complex(phases.cos(), phases.sin()) hv = hv.as_subclass(FHRRTensor) elif self.vsa_tensor == HRRTensor: phases = F.linear(self.bandwidth * input, self.weight) hv = torch.complex(phases.cos(), phases.sin()) hv = torch.real(torch.fft.ifft(torch.fft.ifftshift(hv, dim=1), dim=1)) hv = hv.as_subclass(HRRTensor) return hv