torchjd 0.1.0__tar.gz

torchjd-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Valérian Rey, Pierre Quinton
+
+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.

torchjd-0.1.0/PKG-INFO

@@ -0,0 +1,55 @@
+Metadata-Version: 2.1
+Name: torchjd
+Version: 0.1.0
+Summary: Library for Jacobian Descent with PyTorch.
+Author-email: Valerian Rey <valerian.rey@gmail.com>, Pierre Quinton <pierre.quinton@gmail.com>
+Project-URL: Homepage, https://torchjd.org/
+Project-URL: Documentation, https://torchjd.org/
+Project-URL: Source, https://github.com/TorchJD/torchjd
+Project-URL: Changelog, https://github.com/TorchJD/torchjd/blob/main/CHANGELOG.md
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0.0
+Requires-Dist: quadprog!=0.1.10,>=0.1.9
+Requires-Dist: numpy<2.0.0,>=1.21.0
+Requires-Dist: qpsolvers>=1.0.1
+Requires-Dist: cvxpy>=1.3.0
+
+# ![image](docs/source/icons/favicon-32x32.png) TorchJD
+
+TorchJD is a library enabling Jacobian descent with PyTorch, for optimization of neural networks
+with multiple objectives.
+
+> [!IMPORTANT]
+> This library is currently in an early development stage. The API is subject to significant changes
+> in future versions. Use with caution in production environments and be prepared for potential
+> breaking changes in upcoming releases.
+
+## Installation
+<!-- start installation -->
+TorchJD can be installed directly with pip:
+```bash
+pip install torchjd
+```
+<!-- end installation -->
+
+## Compatibility
+TorchJD requires python 3.10, 3.11 or 3.12. It is only compatible with recent versions of PyTorch
+(>= 2.0). For more information, read the `dependencies` in [pyproject.toml](./pyproject.toml).
+
+## Contribution
+
+Please read the [Contribution page](CONTRIBUTING.md).

torchjd-0.1.0/README.md

@@ -0,0 +1,25 @@
+# ![image](docs/source/icons/favicon-32x32.png) TorchJD
+
+TorchJD is a library enabling Jacobian descent with PyTorch, for optimization of neural networks
+with multiple objectives.
+
+> [!IMPORTANT]
+> This library is currently in an early development stage. The API is subject to significant changes
+> in future versions. Use with caution in production environments and be prepared for potential
+> breaking changes in upcoming releases.
+
+## Installation
+<!-- start installation -->
+TorchJD can be installed directly with pip:
+```bash
+pip install torchjd
+```
+<!-- end installation -->
+
+## Compatibility
+TorchJD requires python 3.10, 3.11 or 3.12. It is only compatible with recent versions of PyTorch
+(>= 2.0). For more information, read the `dependencies` in [pyproject.toml](./pyproject.toml).
+
+## Contribution
+
+Please read the [Contribution page](CONTRIBUTING.md).

torchjd-0.1.0/pyproject.toml

@@ -0,0 +1,61 @@
+[project]
+name = "torchjd"
+version = "0.1.0"
+description = "Library for Jacobian Descent with PyTorch."
+readme = "README.md"
+authors = [
+	{name = "Valerian Rey", email = "valerian.rey@gmail.com"},
+	{name = "Pierre Quinton", email = "pierre.quinton@gmail.com"}
+]
+requires-python = ">=3.10"
+dependencies = [
+    "torch>=2.0.0",
+    "quadprog>=0.1.9, != 0.1.10",  # Doesn't work before 0.1.9, 0.1.10 is yanked
+    "numpy>=1.21.0, <2.0.0",  # Does not work before 1.21. cvxpy is not yet compatible with numpy>=2.0.0. The upper cap should be removed when this becomes the case, or when their pyproject.toml reflects the incompatibility. See https://github.com/cvxpy/cvxpy/issues/2474.
+    "qpsolvers>=1.0.1",  # Does not work before 1.0.1
+    "cvxpy>=1.3.0",  # No Clarabel solver before 1.3.0
+]
+classifiers = [
+	"Development Status :: 4 - Beta",
+	"Intended Audience :: Developers",
+	"Intended Audience :: Science/Research",
+	"License :: OSI Approved :: MIT License",
+	"Operating System :: OS Independent",
+	"Programming Language :: Python",
+	"Programming Language :: Python :: 3",
+	"Programming Language :: Python :: 3.10",
+	"Programming Language :: Python :: 3.11",
+	"Programming Language :: Python :: 3.12",
+	"Topic :: Scientific/Engineering",
+	"Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+[project.urls]
+Homepage = "https://torchjd.org/"
+Documentation = "https://torchjd.org/"
+Source = "https://github.com/TorchJD/torchjd"
+Changelog = "https://github.com/TorchJD/torchjd/blob/main/CHANGELOG.md"
+
+[tool.pdm.dev-dependencies]
+check = [
+	"pre-commit>=2.9.2"  # isort doesn't work before 2.9.2
+]
+
+doc = [
+	"sphinx>=6.0, !=7.2.0, !=7.2.1, !=7.2.3, !=7.2.4, !=7.2.5",  # Versions in [7.2.0, 7.2.5] have a bug with an internal torch import from _C
+	"furo>=2023.0, <2024.04.27",  # Force it to be recent so that the theme looks better, 2024.04.27 seems to have bugged link colors
+	"tomli>=1.1",  # The load function doesn't work similarly before 1.1
+	"sphinx-autodoc-typehints>=1.16.0",  # Some problems with TypeVars before 1.16
+	"myst-parser>=3.0.1"  # Never tested lower versions
+]
+
+test = [
+	"pytest>=7.3",  # Before version 7.3, not all tests are run
+	"contexttimer>=0.3.3, <0.3.4",  # The test requiring contexttimer is not often run, so it could silently break if we uncap this library
+]
+
+plot = [
+    "plotly>=5.19.0",  # Recent version to avoid problems, could be relaxed
+    "dash>=2.16.0",  # Recent version to avoid problems, could be relaxed
+    "kaleido==0.2.1",  # Only works with locked version
+]

torchjd-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

torchjd-0.1.0/src/torchjd/__init__.py

@@ -0,0 +1 @@
+from torchjd.backward import backward

torchjd-0.1.0/src/torchjd/_transform/__init__.py

@@ -0,0 +1,22 @@
+from torchjd._transform.aggregation import make_aggregation
+from torchjd._transform.base import Composition, Conjunction, Transform
+from torchjd._transform.concatenation import Concatenation
+from torchjd._transform.diagonalize import Diagonalize
+from torchjd._transform.grad import Grad
+from torchjd._transform.identity import Identity
+from torchjd._transform.init import Init
+from torchjd._transform.jac import Jac
+from torchjd._transform.matrixify import Matrixify
+from torchjd._transform.reshape import Reshape
+from torchjd._transform.scaling import Scaling
+from torchjd._transform.stack import Stack
+from torchjd._transform.store import Store
+from torchjd._transform.subset import Subset
+from torchjd._transform.tensor_dict import (
+    EmptyTensorDict,
+    Gradients,
+    GradientVectors,
+    JacobianMatrices,
+    Jacobians,
+    TensorDict,
+)

torchjd-0.1.0/src/torchjd/_transform/_differentiation.py

@@ -0,0 +1,34 @@
+from abc import ABC, abstractmethod
+from typing import Iterable, Sequence
+
+from torch import Tensor
+
+from torchjd._transform._utils import ordered_set
+from torchjd._transform.base import _A, Transform
+
+
+class _Differentiation(Transform[_A, _A], ABC):
+    def __init__(self, outputs: Iterable[Tensor], inputs: Iterable[Tensor]):
+        self.outputs = ordered_set(outputs)
+        self.inputs = ordered_set(inputs)
+
+    def _compute(self, tensors: _A) -> _A:
+        tensor_outputs = [tensors[output] for output in self.outputs]
+
+        differentiated_tuple = self._differentiate(tensor_outputs)
+        new_differentiations = dict(zip(self.inputs, differentiated_tuple))
+        return type(tensors)(new_differentiations)
+
+    @abstractmethod
+    def _differentiate(self, tensor_outputs: Sequence[Tensor]) -> tuple[Tensor, ...]:
+        raise NotImplementedError
+
+    @property
+    def required_keys(self) -> set[Tensor]:
+        # outputs in the forward direction become inputs in the backward direction, and vice-versa
+        return set(self.outputs)
+
+    @property
+    def output_keys(self) -> set[Tensor]:
+        # outputs in the forward direction become inputs in the backward direction, and vice-versa
+        return set(self.inputs)

torchjd-0.1.0/src/torchjd/_transform/_utils.py

@@ -0,0 +1,62 @@
+from collections import OrderedDict
+from typing import Hashable, Iterable, Sequence, TypeAlias, TypeVar
+
+import torch
+from torch import Tensor
+
+from torchjd._transform.tensor_dict import EmptyTensorDict, TensorDict, _least_common_ancestor
+
+_KeyType = TypeVar("_KeyType", bound=Hashable)
+_ValueType = TypeVar("_ValueType")
+_OrderedSet: TypeAlias = OrderedDict[_KeyType, None]
+
+_A = TypeVar("_A", bound=TensorDict)
+_B = TypeVar("_B", bound=TensorDict)
+_C = TypeVar("_C", bound=TensorDict)
+
+
+def ordered_set(elements: Iterable[_KeyType]) -> _OrderedSet[_KeyType]:
+    elements = list(elements)
+    result = OrderedDict.fromkeys(elements, None)
+    if len(elements) != len(result):
+        raise ValueError(
+            f"Parameter `elements` should contain unique elements. Found `elements = {elements}`."
+        )
+
+    return result
+
+
+def dicts_union(dicts: Iterable[dict[_KeyType, _ValueType]]) -> dict[_KeyType, _ValueType]:
+    result = {}
+    for d in dicts:
+        result |= d
+    return result
+
+
+def _materialize(
+    optional_tensors: Sequence[Tensor | None], inputs: Sequence[Tensor]
+) -> tuple[Tensor, ...]:
+    """
+    Transforms a sequence of optional tensors by changing each None by a tensor of zeros of the same
+    shape as the corresponding input. Returns the obtained sequence as a tuple.
+
+    Note that the name "materialize" comes from the flag `materialize_grads` from
+    `torch.autograd.grad`, which will be available in future torch releases.
+    """
+
+    tensors = []
+    for optional_tensor, input in zip(optional_tensors, inputs):
+        if optional_tensor is None:
+            tensors.append(torch.zeros_like(input))
+        else:
+            tensors.append(optional_tensor)
+    return tuple(tensors)
+
+
+def _union(tensor_dicts: Iterable[_A]) -> _A:
+    output_type: type[_A] = EmptyTensorDict
+    output: _A = EmptyTensorDict()
+    for tensor_dict in tensor_dicts:
+        output_type = _least_common_ancestor(output_type, type(tensor_dict))
+        output |= tensor_dict
+    return output_type(output)

torchjd-0.1.0/src/torchjd/_transform/aggregation.py

@@ -0,0 +1,12 @@
+from torchjd._transform.base import Transform
+from torchjd._transform.matrixify import Matrixify
+from torchjd._transform.reshape import Reshape
+from torchjd._transform.tensor_dict import Gradients, GradientVectors, JacobianMatrices, Jacobians
+
+
+def make_aggregation(
+    strategy: Transform[JacobianMatrices, GradientVectors]
+) -> Transform[Jacobians, Gradients]:
+    """TODO: doc"""
+
+    return Reshape(strategy.required_keys) << strategy << Matrixify(strategy.required_keys)

torchjd-0.1.0/src/torchjd/_transform/base.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Generic, Sequence
+
+from torch import Tensor
+
+from torchjd._transform._utils import _A, _B, _C, _union
+
+
+class Transform(Generic[_B, _C], ABC):
+    r"""
+    Abstract base class for all transforms. Transforms are elementary building blocks of a jacobian
+    descent backward phase. A transform maps a :class:`~torchjd.transform.tensor_dict.TensorDict` to
+    another. The input :class:`~torchjd.transform.tensor_dict.TensorDict` has keys `required_keys`
+    and the output :class:`~torchjd.transform.tensor_dict.TensorDict` has keys `output_keys`.
+
+    Formally a transform is a function:
+
+    .. math::
+        f:\mathbb R^{n_1+\dots+n_p}\to \mathbb R^{m_1+\dots+m_q}
+
+    where we have ``p`` `required_keys`, ``q`` `output_keys`, ``n_i`` is the number of elements in
+    the value associated to the ``i`` th `required_key` of the input
+    :class:`~torchjd.transform.tensor_dict.TensorDict` and ``m_j`` is the number of elements in the
+    value associated to the ``j`` th `output_key` of the output
+    :class:`~torchjd.transform.tensor_dict.TensorDict`.
+
+    As they are mathematical functions, transforms can be composed together as long as their
+    domains and range meaningfully match [TODO]. We also define conjunction of transforms [TODO].
+    """
+
+    def compose(self, other: Transform[_A, _B]) -> Transform[_A, _C]:
+        return Composition(self, other)
+
+    def conjunct(self, other: Transform[_B, _C]) -> Transform[_B, _C]:
+        return Conjunction([self, other])
+
+    def __repr__(self) -> str:
+        return type(self).__name__
+
+    @abstractmethod
+    def _compute(self, input: _B) -> _C:
+        raise NotImplementedError
+
+    def __call__(self, input: _B) -> _C:
+        input.check_keys_are(self.required_keys)
+        return self._compute(input)
+
+    @property
+    @abstractmethod
+    def required_keys(self) -> set[Tensor]:
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def output_keys(self) -> set[Tensor]:
+        raise NotImplementedError
+
+    __lshift__ = compose
+    __or__ = conjunct
+
+
+class Composition(Transform[_A, _C]):
+    def __init__(self, outer: Transform[_B, _C], inner: Transform[_A, _B]):
+        if outer.required_keys != inner.output_keys:
+            raise ValueError(
+                "The `output_keys` of `inner` must match with the `required_keys` of "
+                f"outer. Found {outer.required_keys} and {inner.output_keys}"
+            )
+        self.outer = outer
+        self.inner = inner
+
+    def __repr__(self) -> str:
+        return repr(self.outer) + " ∘ " + repr(self.inner)
+
+    def _compute(self, input: _A) -> _C:
+        intermediate = self.inner(input)
+        return self.outer(intermediate)
+
+    @property
+    def required_keys(self) -> set[Tensor]:
+        return self.inner.required_keys
+
+    @property
+    def output_keys(self) -> set[Tensor]:
+        return self.outer.output_keys
+
+
+class Conjunction(Transform[_A, _B]):
+    def __init__(self, transforms: Sequence[Transform[_A, _B]]):
+        self.transforms = transforms
+
+        self._required_keys = set(
+            key for transform in transforms for key in transform.required_keys
+        )
+        for transform in transforms:
+            if transform.required_keys != self.required_keys:
+                raise ValueError("All transforms should require the same set of keys.")
+
+        output_keys_with_duplicates = [key for t in transforms for key in t.output_keys]
+        self._output_keys = set(output_keys_with_duplicates)
+
+        if len(self._output_keys) != len(output_keys_with_duplicates):
+            raise ValueError("The sets of output keys of transforms should be disjoint.")
+
+    def _compute(self, tensor_dict: _A) -> _B:
+        output = _union([transform(tensor_dict) for transform in self.transforms])
+        return output
+
+    @property
+    def required_keys(self) -> set[Tensor]:
+        return self._required_keys
+
+    @property
+    def output_keys(self) -> set[Tensor]:
+        return self._output_keys

torchjd-0.1.0/src/torchjd/_transform/concatenation.py

@@ -0,0 +1,65 @@
+from typing import Sequence
+
+import torch
+from torch import Tensor
+
+from torchjd._transform._utils import _A, _materialize, dicts_union
+from torchjd._transform.base import Transform
+from torchjd._transform.tensor_dict import Jacobians
+
+
+class Concatenation(Transform[_A, Jacobians]):
+    def __init__(self, transforms: Sequence[Transform[_A, Jacobians]]):
+        if len(transforms) == 0:
+            raise ValueError("Parameter `transforms` cannot be empty.")
+
+        self.transforms = transforms
+
+        self._required_keys = transforms[0].required_keys
+        self._output_keys = {key for transform in transforms for key in transform.output_keys}
+
+        for transform in transforms[1:]:
+            if transform.required_keys != self.required_keys:
+                raise ValueError("All transforms should require the same set of keys.")
+
+    def _compute(self, input: _A) -> Jacobians:
+        results = [transform(input) for transform in self.transforms]
+        result = _concatenate(results)
+        return result
+
+    @property
+    def required_keys(self) -> set[Tensor]:
+        return self._required_keys
+
+    @property
+    def output_keys(self) -> set[Tensor]:
+        return self._output_keys
+
+
+def _concatenate(jacobians_dicts: list[Jacobians]) -> Jacobians:
+    """
+    Transforms a list of tensor dicts into a single dict of (concatenated) tensors. The set of keys
+    of the resulting dict is the union of the sets of keys of the input dicts. If a key is absent in
+    some input dicts, the corresponding concatenated tensor is filled with zeroes at the positions
+    corresponding to those dicts.
+    """
+
+    # It is important to first remove duplicate keys before computing their associated
+    # concatenated tensor. Otherwise, some computations would be duplicated. Therefore, we first
+    # compute unique_keys, and only then, we compute the concatenated tensors.
+    unique_keys = dicts_union(jacobians_dicts).keys()
+    result = Jacobians({key: _concatenate_one_key(jacobians_dicts, key) for key in unique_keys})
+    return result
+
+
+def _concatenate_one_key(jacobian_dicts: list[Jacobians], input: Tensor) -> Tensor:
+    """
+    Makes the concatenated tensor corresponding to a given key, from a list of tensor dicts.
+    """
+
+    first_dimensions = [jacobian_dict.first_dimension for jacobian_dict in jacobian_dicts]
+    optional_jacobians = [jacobian.get(input, None) for jacobian in jacobian_dicts]
+    expanded_inputs = [input.expand(dim, *input.shape) for dim in first_dimensions]
+    jacobians = _materialize(optional_jacobians, expanded_inputs)
+    jacobian = torch.concatenate(jacobians, dim=0)
+    return jacobian

torchjd-0.1.0/src/torchjd/_transform/diagonalize.py

@@ -0,0 +1,36 @@
+from typing import Iterable
+
+import torch
+from torch import Tensor
+
+from torchjd._transform._utils import ordered_set
+from torchjd._transform.base import Transform
+from torchjd._transform.tensor_dict import Gradients, Jacobians
+
+
+class Diagonalize(Transform[Gradients, Jacobians]):
+    def __init__(self, considered: Iterable[Tensor]):
+        self.considered = ordered_set(considered)
+        self.indices: list[tuple[int, int]] = []
+        begin = 0
+        for tensor in self.considered:
+            end = begin + tensor.numel()
+            self.indices.append((begin, end))
+            begin = end
+
+    def _compute(self, tensors: Gradients) -> Jacobians:
+        flattened_considered_values = [tensors[key].reshape([-1]) for key in self.considered]
+        diagonal_matrix = torch.cat(flattened_considered_values).diag()
+        diagonalized_tensors = {
+            key: diagonal_matrix[:, begin:end].reshape((-1,) + key.shape)
+            for (begin, end), key in zip(self.indices, self.considered)
+        }
+        return Jacobians(diagonalized_tensors)
+
+    @property
+    def required_keys(self) -> set[Tensor]:
+        return set(self.considered)
+
+    @property
+    def output_keys(self) -> set[Tensor]:
+        return set(self.considered)

torchjd-0.1.0/src/torchjd/_transform/grad.py

@@ -0,0 +1,60 @@
+from typing import Iterable, Sequence
+
+import torch
+from torch import Tensor
+
+from torchjd._transform._differentiation import _Differentiation
+from torchjd._transform._utils import _materialize
+from torchjd._transform.tensor_dict import Gradients
+
+
+class Grad(_Differentiation[Gradients]):
+    def __init__(
+        self,
+        outputs: Iterable[Tensor],
+        inputs: Iterable[Tensor],
+        retain_graph: bool = False,
+    ):
+        super().__init__(outputs, inputs)
+        self.retain_graph = retain_graph
+
+    def _differentiate(self, grad_outputs: Sequence[Tensor]) -> tuple[Tensor, ...]:
+        return _grad(
+            outputs=self.outputs,
+            inputs=self.inputs,
+            grad_outputs=grad_outputs,
+            retain_graph=self.retain_graph,
+            create_graph=False,
+            allow_unused=True,
+        )
+
+
+def _grad(
+    outputs: Sequence[Tensor],
+    inputs: Sequence[Tensor],
+    grad_outputs: Sequence[Tensor],
+    retain_graph: bool,
+    create_graph: bool,
+    allow_unused: bool,
+) -> tuple[Tensor, ...]:
+    """
+    Wraps `autograd.grad` to give it additional responsibilities that it should have (like being
+    able to work with an empty sequence of `inputs`).
+    """
+
+    if len(inputs) == 0:
+        return tuple()
+
+    if len(outputs) == 0:
+        return tuple([torch.empty(input.shape) for input in inputs])
+
+    optional_grads = torch.autograd.grad(
+        outputs,
+        inputs,
+        grad_outputs=grad_outputs,
+        retain_graph=retain_graph,
+        create_graph=create_graph,
+        allow_unused=allow_unused,
+    )
+    grads = _materialize(optional_grads, inputs)
+    return grads

torchjd-0.1.0/src/torchjd/_transform/identity.py

@@ -0,0 +1,22 @@
+from typing import Iterable
+
+from torch import Tensor
+
+from torchjd._transform._utils import _A
+from torchjd._transform.base import Transform
+
+
+class Identity(Transform[_A, _A]):
+    def __init__(self, required_keys: Iterable[Tensor]):
+        self._required_keys = set(required_keys)
+
+    def _compute(self, tensor_dict: _A) -> _A:
+        return tensor_dict
+
+    @property
+    def required_keys(self) -> set[Tensor]:
+        return self._required_keys
+
+    @property
+    def output_keys(self) -> set[Tensor]:
+        return self._required_keys

torchjd-0.1.0/src/torchjd/_transform/init.py

@@ -0,0 +1,30 @@
+from typing import Iterable
+
+import torch
+from torch import Tensor
+
+from torchjd._transform.base import Transform
+from torchjd._transform.tensor_dict import EmptyTensorDict, Gradients
+
+
+class Init(Transform[EmptyTensorDict, Gradients]):
+    def __init__(self, values: Iterable[Tensor]):
+        self.values = set(values)
+
+    def _compute(self, input: EmptyTensorDict) -> Gradients:
+        r"""
+        Computes the gradients of the ``value`` with respect to itself. Returns the result as a
+        dictionary. The only key of the dictionary is ``value``. The corresponding gradient is a
+        tensor of 1s of identical shape, because :math:`\frac{\partial v}{\partial v} = 1` for any
+        :math:`v`.
+        """
+
+        return Gradients({value: torch.ones_like(value) for value in self.values})
+
+    @property
+    def required_keys(self) -> set[Tensor]:
+        return set()
+
+    @property
+    def output_keys(self) -> set[Tensor]:
+        return self.values