torchjd 0.1.0__py3-none-any.whl

torchjd/__init__.py

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

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/_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/_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/_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/_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/_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/_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/_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/_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/_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

torchjd/_transform/jac.py

@@ -0,0 +1,109 @@
+from itertools import accumulate
+from typing import Iterable, Sequence
+
+import torch
+from torch import Size, Tensor
+
+from torchjd._transform._differentiation import _Differentiation
+from torchjd._transform._utils import _materialize
+from torchjd._transform.tensor_dict import Jacobians
+
+
+class Jac(_Differentiation[Jacobians]):
+    def __init__(
+        self,
+        outputs: Iterable[Tensor],
+        inputs: Iterable[Tensor],
+        chunk_size: int | None,
+        retain_graph: bool = False,
+    ):
+        super().__init__(outputs, inputs)
+        self.chunk_size = chunk_size
+        self.retain_graph = retain_graph
+
+    def _differentiate(self, jac_outputs: Sequence[Tensor]) -> tuple[Tensor, ...]:
+        return _jac(
+            outputs=self.outputs,
+            inputs=self.inputs,
+            jac_outputs=jac_outputs,
+            chunk_size=self.chunk_size,
+            retain_graph=self.retain_graph,
+            create_graph=False,
+            allow_unused=True,
+        )
+
+
+def _jac(
+    outputs: Sequence[Tensor],
+    inputs: Sequence[Tensor],
+    jac_outputs: Sequence[Tensor],
+    chunk_size: int | None,
+    retain_graph: bool,
+    create_graph: bool,
+    allow_unused: bool,
+) -> tuple[Tensor, ...]:
+    """
+    Wraps the call to `autograd.grad` to compute the jacobian with respect to each input, in an
+    optimized way. The first dimension of the jacobians is equal to the length of the sequence of
+    `outputs`, which should be the same as the length of the sequence of `grad_outputs`. This should
+    be equivalent to calling `_grad(outputs[i], inputs, grad_outputs[i], ...)` for all i
+    sequentially, and stacking the elements of each resulting tuple.
+    """
+
+    if len(inputs) == 0:
+        return tuple()
+
+    n_outputs = len(outputs)
+    if len(jac_outputs) != n_outputs:
+        raise ValueError(
+            "Parameters `outputs` and `jac_outputs` should be sequences of the same length. Found "
+            f"`len(outputs) = {n_outputs}` and `len(jac_outputs) = {len(jac_outputs)}`."
+        )
+
+    if n_outputs == 0:
+        return tuple([torch.empty((0,) + input.shape) for input in inputs])
+
+    def get_vjp(v):
+        optional_grads = torch.autograd.grad(
+            outputs,
+            inputs,
+            grad_outputs=v,
+            retain_graph=retain_graph,
+            create_graph=create_graph,
+            allow_unused=allow_unused,
+        )
+        grads = _materialize(optional_grads, inputs=inputs)
+        return torch.concatenate([grad.reshape([-1]) for grad in grads])
+
+    grouped_jacobian_matrix = torch.vmap(get_vjp, chunk_size=chunk_size)(jac_outputs)
+
+    lengths = [input.numel() for input in inputs]
+    jacobian_matrices = _extract_sub_matrices(grouped_jacobian_matrix, lengths)
+
+    shapes = [input.shape for input in inputs]
+    jacobians = _reshape_matrices(jacobian_matrices, shapes)
+
+    return tuple(jacobians)
+
+
+def _extract_sub_matrices(matrix: Tensor, lengths: Sequence[int]) -> list[Tensor]:
+    cumulative_lengths = [*accumulate(lengths)]
+
+    if cumulative_lengths[-1] != matrix.shape[1]:
+        raise ValueError(
+            "The sum of the provided lengths should be equal to the number of columns in the "
+            "provided matrix."
+        )
+
+    start_indices = [0] + cumulative_lengths[:-1]
+    end_indices = cumulative_lengths
+    return [matrix[:, start:end] for start, end in zip(start_indices, end_indices)]
+
+
+def _reshape_matrices(matrices: Sequence[Tensor], shapes: Sequence[Size]) -> Sequence[Tensor]:
+    if len(matrices) != len(shapes):
+        raise ValueError(
+            "Parameters `matrices` and `shapes` should contain the same number of elements."
+        )
+
+    return [matrix.view((matrix.shape[0],) + shape) for matrix, shape in zip(matrices, shapes)]

torchjd/_transform/matrixify.py

@@ -0,0 +1,25 @@
+from typing import Iterable
+
+from torch import Tensor
+
+from torchjd._transform.base import Transform
+from torchjd._transform.tensor_dict import JacobianMatrices, Jacobians
+
+
+class Matrixify(Transform[Jacobians, JacobianMatrices]):
+    def __init__(self, required_keys: Iterable[Tensor]):
+        self._required_keys = set(required_keys)
+
+    def _compute(self, jacobians: Jacobians) -> JacobianMatrices:
+        jacobian_matrices = {
+            key: jacobian.view(jacobian.shape[0], -1) for key, jacobian in jacobians.items()
+        }
+        return JacobianMatrices(jacobian_matrices)
+
+    @property
+    def required_keys(self) -> set[Tensor]:
+        return self._required_keys
+
+    @property
+    def output_keys(self) -> set[Tensor]:
+        return self._required_keys

torchjd/_transform/reshape.py

@@ -0,0 +1,26 @@
+from typing import Iterable
+
+from torch import Tensor
+
+from torchjd._transform.base import Transform
+from torchjd._transform.tensor_dict import Gradients, GradientVectors
+
+
+class Reshape(Transform[GradientVectors, Gradients]):
+    def __init__(self, required_keys: Iterable[Tensor]):
+        self._required_keys = set(required_keys)
+
+    def _compute(self, gradient_vectors: GradientVectors) -> Gradients:
+        gradients = {
+            key: gradient_vector.view(key.shape)
+            for key, gradient_vector in gradient_vectors.items()
+        }
+        return Gradients(gradients)
+
+    @property
+    def required_keys(self) -> set[Tensor]:
+        return self._required_keys
+
+    @property
+    def output_keys(self) -> set[Tensor]:
+        return self._required_keys

torchjd/_transform/scaling.py

@@ -0,0 +1,21 @@
+from torch import Tensor
+
+from torchjd._transform._utils import _A
+from torchjd._transform.base import Transform
+
+
+class Scaling(Transform[_A, _A]):
+    def __init__(self, scalings: dict[Tensor, float]):
+        self.scalings = scalings
+
+    def _compute(self, tensor_dict: _A) -> _A:
+        output = {key: scaling * tensor_dict[key] for key, scaling in self.scalings.items()}
+        return type(tensor_dict)(output)
+
+    @property
+    def required_keys(self) -> set[Tensor]:
+        return set(self.scalings.keys())
+
+    @property
+    def output_keys(self) -> set[Tensor]:
+        return set(self.scalings.keys())