From 8b8193f1324dea233237a9362d0f508c6c533d10 Mon Sep 17 00:00:00 2001 From: lezcano Date: Sat, 6 Jun 2026 20:47:44 +0100 Subject: [PATCH 1/6] Modernize for PyTorch 2.6 --- README.rst | 8 +- docs/README.md | 4 +- examples/copying_problem.py | 14 +- examples/eigenvalue.py | 13 +- examples/parametrisations.ipynb | 20 +- examples/sequential_mnist.py | 18 +- geotorch/__init__.py | 1 - geotorch/almostorthogonal.py | 5 +- geotorch/constraints.py | 24 +- geotorch/fixedrank.py | 8 +- geotorch/hurwitz.py | 44 +-- geotorch/lowrank.py | 2 +- geotorch/parametrize.py | 670 -------------------------------- geotorch/pssdfixedrank.py | 2 +- geotorch/skew.py | 5 +- geotorch/sl.py | 47 ++- geotorch/so.py | 7 +- geotorch/sphere.py | 39 +- geotorch/symmetric.py | 2 +- geotorch/utils.py | 2 +- setup.cfg | 3 - test/test_hurwitz.py | 81 ++-- test/test_integration.py | 10 +- test/test_lowrank.py | 15 +- test/test_skew.py | 10 +- test/test_sl.py | 25 ++ test/test_sphere.py | 36 ++ test/test_symmetric.py | 8 +- 28 files changed, 278 insertions(+), 845 deletions(-) delete mode 100644 geotorch/parametrize.py diff --git a/README.rst b/README.rst index 024ceb87..712e33cc 100644 --- a/README.rst +++ b/README.rst @@ -147,12 +147,14 @@ If one wants to use a parametrized tensor in different places in their model, or .. code:: python - with geotorch.parametrize.cached(): + from torch.nn.utils import parametrize + + with parametrize.cached(): logits = model(input_) Of course, this ``with`` statement may be used simply inside the forward function where the parametrized layer is used several times. -These ideas fall in the context of parametrized optimization, where one wraps a tensor ``X`` with a function ``f``, and rather than using ``X``, uses ``f(X)``. Particular examples of this idea are pruning, weight normalization, and spectral normalization among others. This repository implements a framework to approach this kind of problems. This framework was accepted to core PyTorch 1.8. It can be found under `torch.nn.utils.parametrize`_ and `torch.nn.utils.parametrizations`_. When using PyTorch 1.10 or higher, the native PyTorch functions are used within GeoTorch. In this case, the user can interact with the parametrizations in GeoTorch using the PyTorch functions. +These ideas fall in the context of parametrized optimization, where one wraps a tensor ``X`` with a function ``f``, and rather than using ``X``, uses ``f(X)``. Particular examples of this idea are pruning, weight normalization, and spectral normalization among others. As every space in GeoTorch is, at its core, a map from a flat space into a manifold, the tools implemented here also serve as a building block in normalizing flows. Using a factorized space such as |low|_ it is direct to compute the determinant of the transformation it defines, as we have direct access to the singular values of the layer. @@ -226,8 +228,6 @@ Please cite the following work if you found GeoTorch useful. This paper exposes .. _here: https://github.com/Lezcano/geotorch/blob/master/examples/copying_problem.py#L16 .. _torch.nn.utils.parametrize: https://pytorch.org/docs/stable/generated/torch.nn.utils.parametrize.register_parametrization.html .. _torch.nn.utils.parametrizations: https://pytorch.org/docs/stable/generated/torch.nn.utils.parametrizations.orthogonal.html -.. _geotorch/parametrize.py: https://github.com/Lezcano/geotorch/blob/master/geotorch/parametrize.py .. _examples/sequential_mnist.py: https://github.com/Lezcano/geotorch/blob/master/examples/sequential_mnist.py .. _examples/copying_problem.py: https://github.com/Lezcano/geotorch/blob/master/examples/copying_problem.py .. _examples/parametrisations.ipynb: https://github.com/Lezcano/geotorch/blob/master/examples/parametrisations.ipynb - diff --git a/docs/README.md b/docs/README.md index a8391ceb..a39618ff 100644 --- a/docs/README.md +++ b/docs/README.md @@ -6,10 +6,10 @@ To make the documentation you can run make html ``` -To generate the docs from the GeoTorch source excluding the `parametrize.py` file, run +To generate the docs from the GeoTorch source, run ``` -SPHINX_APIDOC_OPTIONS=members sphinx-apidoc -o ./source ../geotorch ../geotorch/parametrize.py +SPHINX_APIDOC_OPTIONS=members sphinx-apidoc -o ./source ../geotorch ``` To check the spelling diff --git a/examples/copying_problem.py b/examples/copying_problem.py index 185ab5a1..4fef12f8 100644 --- a/examples/copying_problem.py +++ b/examples/copying_problem.py @@ -20,6 +20,7 @@ import torch from torch import nn import torch.nn.functional as F +from torch.nn.utils import parametrize import geotorch @@ -47,11 +48,12 @@ class modrelu(nn.Module): def __init__(self, features): super(modrelu, self).__init__() self.features = features - self.b = nn.Parameter(torch.Tensor(self.features)) + self.b = nn.Parameter(torch.empty(self.features)) self.reset_parameters() def reset_parameters(self): - self.b.data.uniform_(-0.01, 0.01) + with torch.no_grad(): + self.b.uniform_(-0.01, 0.01) def forward(self, inputs): norm = torch.abs(inputs) @@ -77,7 +79,7 @@ def __init__(self, input_size, hidden_size): self.reset_parameters() def reset_parameters(self): - nn.init.kaiming_normal_(self.input_kernel.weight.data, nonlinearity="relu") + nn.init.kaiming_normal_(self.input_kernel.weight, nonlinearity="relu") # The manifold class is under `layer.parametrizations.tensor_name[0]` M = self.recurrent_kernel.parametrizations.weight[0] # Every manifold has a convenience sample method, but you can use your own initializer @@ -103,13 +105,13 @@ def __init__(self, alphabet_size, hidden_size): self.reset_parameters() def reset_parameters(self): - nn.init.kaiming_normal_(self.lin.weight.data, nonlinearity="relu") - nn.init.constant_(self.lin.bias.data, 0) + nn.init.kaiming_normal_(self.lin.weight, nonlinearity="relu") + nn.init.constant_(self.lin.bias, 0) def forward(self, inputs): out_rnn = self.rnn.default_hidden(inputs[:, 0, ...]) outputs = [] - with geotorch.parametrize.cached(): + with parametrize.cached(): for input in torch.unbind(inputs, dim=1): out_rnn = self.rnn(input, out_rnn) outputs.append(self.lin(out_rnn)) diff --git a/examples/eigenvalue.py b/examples/eigenvalue.py index 8e25e17e..3cb6cc50 100644 --- a/examples/eigenvalue.py +++ b/examples/eigenvalue.py @@ -6,15 +6,6 @@ import torch -try: - from torch.linalg import eigvalsh -except ImportError: - from torch import symeig - - def eigvalsh(X): - return symeig(X, eigenvectors=False).eigenvalues - - from torch import nn import geotorch @@ -33,7 +24,7 @@ def __init__(self, n): def forward(self, A): x = self.x - return x.T @ A @ x + return x @ A @ x # Generate matrix @@ -41,7 +32,7 @@ def forward(self, A): A = 0.5 * (A + A.T) # Compare against diagonalization (eigenvalues are returend in ascending order) -max_eigenvalue = eigvalsh(A)[-1] +max_eigenvalue = torch.linalg.eigvalsh(A)[-1] print("Max eigenvalue: {:10.5f}".format(max_eigenvalue)) # Instantiate model and optimiser diff --git a/examples/parametrisations.ipynb b/examples/parametrisations.ipynb index 82859d98..06141c82 100644 --- a/examples/parametrisations.ipynb +++ b/examples/parametrisations.ipynb @@ -82,7 +82,7 @@ "\n", "```python\n", "layer = nn.Linear(N, N)\n", - "torch.register_parametrization(layer, \"weight\", SymmetricParametrization())\n", + "torch.nn.utils.parametrize.register_parametrization(layer, \"weight\", SymmetricParametrization())\n", "# layer now behaves as an object from the `Symmetric` class\n", "print(layer.weight) # Prints the symmetric matrix\n", "layer(x) # Multiplies the vectors `x` by the symmetric matrix layer.weight \n", @@ -127,16 +127,16 @@ "source": [ "We would like to use it as:\n", "```python\n", - "cnn = nn.Conv2D(8, 16, (3, 3))\n", - "torch.register_parametrization(cnn, \"weight\", PruningParametrization(cnn.weight, p_drop=0.1))\n", + "cnn = nn.Conv2d(8, 16, (3, 3))\n", + "torch.nn.utils.parametrize.register_parametrization(cnn, \"weight\", PruningParametrization(cnn.weight, p_drop=0.1))\n", "# 10% of the entires of the tensor cnn.weight have now been zeroed out\n", "```\n", "### Other examples:\n", - "- `torch.weight_norm`\n", - "- `torch.spectral_norm` (to regularise the Lipschitz constant of a layer)\n", + "- `torch.nn.utils.parametrizations.weight_norm`\n", + "- `torch.nn.utils.parametrizations.spectral_norm` (to regularise the Lipschitz constant of a layer)\n", "- Optimisation with orthogonal constraints / invertible layers / Symmetric Positive Defininite layers... More on this later\n", "\n", - "## Implementing `torch.register_parametrization`\n", + "## Implementing `torch.nn.utils.parametrize.register_parametrization`\n", "### A first approximation\n", "A moment's reflection shows that it is possible to implement `Symmetric` without having to reimplement `nn.Linear` by using inheritance and properties.\n", "\n", @@ -279,7 +279,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "## Other things that `torch.register_parametrization` implements:\n", + "## Other things that `torch.nn.utils.parametrize.register_parametrization` implements:\n", "\n", "- If the module implements a `right_inverse` method (similar to a right-inverse of forward, more on this below), it allows initialising the parametrised buffer/parameter\n", "- It allows putting several parametrisations on the same buffer/parameter\n", @@ -302,9 +302,7 @@ "metadata": {}, "outputs": [], "source": [ - "# This part assumes that you have `geotorch` installed. You can install it doing\n", - "# pip install git+https://github.com/Lezcano/geotorch/\n", - "import geotorch.parametrize as P\n", + "import torch.nn.utils.parametrize as P\n", "\n", "class Skew(nn.Module):\n", " def forward(self, X):\n", @@ -350,7 +348,7 @@ " # Cayley map: (I + A)(I - A)^{-1}\n", " # This is orthogonal whenever A is skew-symmetric\n", " Id = torch.eye(A.size(0))\n", - " return self.B @ torch.solve(Id - A, Id + A).solution\n", + " return self.B @ torch.linalg.solve(Id - A, Id + A)\n", "\n", " def is_orthogonal(self, X):\n", " Id = torch.eye(X.size(0))\n", diff --git a/examples/sequential_mnist.py b/examples/sequential_mnist.py index 30ec473d..77989654 100644 --- a/examples/sequential_mnist.py +++ b/examples/sequential_mnist.py @@ -13,6 +13,7 @@ import torch import torch.nn as nn +from torch.nn.utils import parametrize import math import argparse from torchvision import datasets, transforms @@ -55,11 +56,12 @@ class modrelu(nn.Module): def __init__(self, features): super(modrelu, self).__init__() self.features = features - self.b = nn.Parameter(torch.Tensor(self.features)) + self.b = nn.Parameter(torch.empty(self.features)) self.reset_parameters() def reset_parameters(self): - self.b.data.uniform_(-0.01, 0.01) + with torch.no_grad(): + self.b.uniform_(-0.01, 0.01) def forward(self, inputs): norm = torch.abs(inputs) @@ -92,15 +94,15 @@ def __init__(self, input_size, hidden_size): self.reset_parameters() def reset_parameters(self): - nn.init.kaiming_normal_(self.input_kernel.weight.data, nonlinearity="relu") + nn.init.kaiming_normal_(self.input_kernel.weight, nonlinearity="relu") # Initialize the recurrent kernel à la Cayley, as having a block-diagonal matrix # seems to help in classification problems def init_(x): x.uniform_(0.0, math.pi / 2.0) - c = torch.cos(x.data) - x.data = -torch.sqrt((1.0 - c) / (1.0 + c)) + c = torch.cos(x) + x.copy_(-torch.sqrt((1.0 - c) / (1.0 + c))) K = self.recurrent_kernel # We initialize it by assigning directly to it from a sampler @@ -130,7 +132,7 @@ def forward(self, inputs): if self.permute: inputs = inputs[:, self.permutation] out_rnn = self.rnn.default_hidden(inputs[:, 0, ...]) - with geotorch.parametrize.cached(): + with parametrize.cached(): for input in torch.unbind(inputs, dim=1): out_rnn = self.rnn(input.unsqueeze(dim=1), out_rnn) return self.lin(out_rnn) @@ -154,11 +156,11 @@ def main(): datasets.MNIST( "./mnist", train=True, download=True, transform=transforms.ToTensor() ), - **kwargs + **kwargs, ) test_loader = torch.utils.data.DataLoader( datasets.MNIST("./mnist", train=False, transform=transforms.ToTensor()), - **kwargs + **kwargs, ) # Model and optimizers diff --git a/geotorch/__init__.py b/geotorch/__init__.py index ed3a3384..05d2dfb7 100644 --- a/geotorch/__init__.py +++ b/geotorch/__init__.py @@ -35,7 +35,6 @@ from .hurwitz import Hurwitz from .utils import update_base - __version__ = "0.3.0" diff --git a/geotorch/almostorthogonal.py b/geotorch/almostorthogonal.py index 80d978a8..f38720b5 100644 --- a/geotorch/almostorthogonal.py +++ b/geotorch/almostorthogonal.py @@ -9,8 +9,7 @@ def scaled_sigmoid(t): def inv_scaled_sigmoid(t): - y = 0.5 * t + 0.5 - return torch.log(y / (1.0 - y)) + return torch.logit(0.5 * t + 0.5) class AlmostOrthogonal(LowRank): @@ -131,7 +130,7 @@ def sample(self, distribution="uniform", init_=None): >>> layer = nn.Linear(20, 20) >>> M = AlmostOrthogonal(layer.weight.size(), lam=0.5) - >>> geotorch.register_parametrization(layer, "weight", M) + >>> torch.nn.utils.parametrize.register_parametrization(layer, "weight", M) >>> layer.weight = M.sample() Args: diff --git a/geotorch/constraints.py b/geotorch/constraints.py index f1054b06..cc8ff422 100644 --- a/geotorch/constraints.py +++ b/geotorch/constraints.py @@ -1,5 +1,5 @@ import torch -import geotorch.parametrize as P +from torch.nn.utils import parametrize from .symmetric import Symmetric from .skew import Skew @@ -24,13 +24,13 @@ def _register_manifold(module, tensor_name, cls, *args): # Initialize without checking in manifold X = M.sample() - if not P.is_parametrized(module, tensor_name): + if not parametrize.is_parametrized(module, tensor_name): with torch.no_grad(): tensor.copy_(X) else: setattr(module, tensor_name, X) - P.register_parametrization(module, tensor_name, M, unsafe=True) + parametrize.register_parametrization(module, tensor_name, M, unsafe=True) return module @@ -58,7 +58,7 @@ def symmetric(module, tensor_name="weight", lower=True): lower (bool): Optional. Uses the lower triangular part of the matrix to parametrize the matrix. Default: ``True`` """ - P.register_parametrization(module, tensor_name, Symmetric(lower)) + parametrize.register_parametrization(module, tensor_name, Symmetric(lower)) return module @@ -101,10 +101,10 @@ def sphere(module, tensor_name="weight", radius=1.0, embedded=False): >>> layer = nn.Linear(20, 30) >>> geotorch.sphere(layer, "bias") - >>> torch.norm(layer.bias) + >>> torch.linalg.vector_norm(layer.bias) tensor(1.) >>> geotorch.sphere(layer, "weight") # Make the columns unit norm - >>> torch.allclose(torch.norm(layer.weight, dim=-1), torch.ones(30)) + >>> torch.allclose(torch.linalg.vector_norm(layer.weight, dim=-1), torch.ones(30)) True Args: @@ -135,12 +135,12 @@ def orthogonal(module, tensor_name="weight", triv="expm"): >>> layer = nn.Linear(20, 30) >>> geotorch.orthogonal(layer, "weight") - >>> torch.norm(layer.weight.T @ layer.weight - torch.eye(20,20)) + >>> torch.linalg.matrix_norm(layer.weight.T @ layer.weight - torch.eye(20,20)) tensor(4.8488e-05) >>> layer = nn.Conv2d(20, 40, 3, 3) # Make the kernels orthogonal >>> geotorch.orthogonal(layer, "weight") - >>> torch.norm(layer.weight.transpose(-2, -1) @ layer.weight - torch.eye(3,3)) + >>> torch.linalg.matrix_norm(layer.weight.transpose(-2, -1) @ layer.weight - torch.eye(3,3)) tensor(1.2225e-05) Args: @@ -221,12 +221,12 @@ def grassmannian(module, tensor_name="weight", triv="expm"): >>> layer = nn.Linear(20, 30) >>> geotorch.grassmannian(layer, "weight") - >>> torch.norm(layer.weight.t() @ layer.weight - torch.eye(20,20)) + >>> torch.linalg.matrix_norm(layer.weight.t() @ layer.weight - torch.eye(20,20)) tensor(1.8933e-05) >>> layer = nn.Conv2d(20, 40, 3, 3) # Make the kernels represent subspaces >>> geotorch.grassmannian(layer, "weight") - >>> torch.norm(layer.weight.transpose(-2, -1) @ layer.weight - torch.eye(3,3)) + >>> torch.linalg.matrix_norm(layer.weight.transpose(-2, -1) @ layer.weight - torch.eye(3,3)) tensor(8.3796-06) Args: @@ -566,8 +566,8 @@ def alpha_stable(module: torch.nn.Module, tensor_name="weight", alpha=1e-3): >>> layer = nn.Linear(30, 30) >>> alpha_stable(layer, "weight", alpha=2) - >>> torch.all(torch.real(torch.linalg.eigvals(layer.weight))<=-layer.alpha) - True + >>> torch.all(torch.linalg.eigvals(layer.weight).real <= -2) + tensor(True) Args: module (nn.Module): module on which to register the parametrization diff --git a/geotorch/fixedrank.py b/geotorch/fixedrank.py index 36547f1a..7b71d780 100644 --- a/geotorch/fixedrank.py +++ b/geotorch/fixedrank.py @@ -9,7 +9,7 @@ def softplus_epsilon(x, epsilon=1e-6): def inv_softplus_epsilon(x, epsilon=1e-6): y = x - epsilon - return torch.where(y > 20, y, y.expm1().log()) + return y + (-y).expm1().neg().log() class FixedRank(LowRank): @@ -84,10 +84,8 @@ def in_manifold_singular_values(self, S, eps=1e-5): """ if not super().in_manifold_singular_values(S, eps): return False - # We compute the \infty-norm of the eigenvalues D = S[..., : self.rank] - infty_norm = D.abs().max(dim=-1).values - return (infty_norm > eps).all().item() + return (D > eps).all().item() def sample(self, init_=torch.nn.init.xavier_normal_, eps=5e-6, factorized=False): r""" @@ -102,7 +100,7 @@ def sample(self, init_=torch.nn.init.xavier_normal_, eps=5e-6, factorized=False) >>> layer = nn.Linear(20, 20) >>> M = FixedRank(layer.weight.size(), rank=6) - >>> geotorch.register_parametrization(layer, "weight", M) + >>> torch.nn.utils.parametrize.register_parametrization(layer, "weight", M) >>> layer.weight = M.sample() Args: diff --git a/geotorch/hurwitz.py b/geotorch/hurwitz.py index 63fbf1a0..ebd5f3a4 100644 --- a/geotorch/hurwitz.py +++ b/geotorch/hurwitz.py @@ -12,7 +12,7 @@ def get_lyap_exp(A): - return -torch.max(torch.real(torch.linalg.eigvals(A))) + return -torch.linalg.eigvals(A).real.max() class Hurwitz(ProductManifold): @@ -38,17 +38,15 @@ def __init__(self, size, alpha: float = 0.0, triv="expm"): callable. Default: ``"expm"`` """ - assert alpha >= 0, ValueError(f"alpha must be positive found {alpha}") + if alpha < 0: + raise ValueError(f"alpha must be non-negative, found {alpha}") n, tensorial_size = Hurwitz.parse_size(size) super().__init__(Hurwitz.manifolds(n, tensorial_size, triv)) self.n = n self.tensorial_size = tensorial_size - self.alpha = alpha self.register_buffer("In", torch.eye(n).expand(*self.tensorial_size, n, n)) - self.register_buffer( - "alphaIn", self.alpha * torch.eye(n).expand(*self.tensorial_size, n, n) - ) + self.register_buffer("alpha", torch.as_tensor(alpha)) @classmethod def parse_size(cls, size): @@ -66,7 +64,7 @@ def manifolds(n, tensorial_size, triv): return PSD(size, triv=triv), PSD(size, triv=triv), Skew(size) def submersion(self, Q, P, S): - return P @ torch.add(S, Q, alpha=-0.5) - self.alphaIn + return P @ torch.add(S, Q, alpha=-0.5) - self.alpha * self.In def forward(self, X1, X2, X3): Q, P, S = super().forward([X1, X2, X3]) @@ -85,11 +83,10 @@ def submersion_inv(self, A: torch.Tensor, check_in_manifold=True, rho=1, tol=1e- """ if check_in_manifold and not self.in_manifold_eigen(A): - print(f"Enforced alpha = {self.alpha} and given alpha = {get_lyap_exp(A)}") - raise InManifoldError(get_lyap_exp(A), self.alpha) + raise InManifoldError(A, self) with torch.no_grad(): - A_shifted = A + (self.alpha - tol) * self.In + A_shifted = A + self.alpha * self.In A_shifted_T = A_shifted.mT.contiguous() # Batched Kronecker product using vmap for cleaner implementation @@ -109,15 +106,18 @@ def _single_kron_sum(A_single): vec_P = torch.linalg.solve(M, -flat_Q.unsqueeze(-1)).squeeze(-1) P = vec_P.view(*self.tensorial_size, self.n, self.n) - residual = torch.norm(A.mT @ P + P @ A + 2 * (self.alpha - tol) * P + Q) + residual = torch.linalg.vector_norm( + A.mT @ P + P @ A + 2 * self.alpha * P + Q + ) if residual >= tol: raise ValueError( - f"Lyapunov equation ill-conditioned solve failed. \n Residual norm: {torch.norm(residual):.2e}" + "Lyapunov equation ill-conditioned solve failed. " + f"Residual norm: {residual:.2e}" ) - S = P @ (A + (self.alpha - tol) * self.In) + 0.5 * Q + S = P @ A_shifted + 0.5 * Q - P_inv = torch.inverse(P) + P_inv = torch.linalg.inv(P) return Q, P_inv, S @@ -140,13 +140,14 @@ def sample(self, init_=torch.nn.init.xavier_normal_): .. math:: A = P^{-1}(-Q/2 + S) - \alpha I - with :math:`P^{-1}_{i,j}, Q_{i,j}, S_{i,j} \sim \texttt{init_}` + The positive-definite factors are shifted by the identity to keep the + sample away from the boundary of the Hurwitz matrices. The output of this method can be used to initialize a parametrized tensor as:: >>> layer = nn.Linear(20, 20) >>> M = Hurwitz(layer.weight.size(), alpha=0.5) - >>> geotorch.register_parametrization(layer, "weight", M) + >>> torch.nn.utils.parametrize.register_parametrization(layer, "weight", M) >>> layer.weight = M.sample() Args: @@ -154,17 +155,12 @@ def sample(self, init_=torch.nn.init.xavier_normal_): Default: ``torch.nn.init.xavier_normal_`` """ with torch.no_grad(): - - mani_psd = PSD((self.n, self.n)) - mani_skew = Skew((self.n, self.n)) - - P = mani_psd.sample(init_) - Q = mani_psd.sample(init_) - S = mani_skew.sample(init_) + P = self[0].sample(init_) + self.In + Q = self[1].sample(init_) + self.In + S = self[2](init_(torch.empty_like(P))) A = self.submersion(Q, P, S) return A def extra_repr(self) -> str: - print(self.alpha) return _extra_repr(n=self.n, alpha=self.alpha) diff --git a/geotorch/lowrank.py b/geotorch/lowrank.py index 6664cb16..90c09cd0 100644 --- a/geotorch/lowrank.py +++ b/geotorch/lowrank.py @@ -137,7 +137,7 @@ def sample(self, init_=torch.nn.init.xavier_normal_, factorized=False): >>> layer = nn.Linear(20, 20) >>> M = LowRank(layer.weight.size(), rank=6) - >>> geotorch.register_parametrization(layer, "weight", M) + >>> torch.nn.utils.parametrize.register_parametrization(layer, "weight", M) >>> layer.weight = M.sample() Args: diff --git a/geotorch/parametrize.py b/geotorch/parametrize.py deleted file mode 100644 index 313163b0..00000000 --- a/geotorch/parametrize.py +++ /dev/null @@ -1,670 +0,0 @@ -import torch - -if int(torch.__version__.split(".")[1]) >= 10: - from torch.nn.utils.parametrize import * -else: - from torch.nn.modules.container import ModuleList, ModuleDict, Module - from torch.nn.parameter import Parameter - from torch import Tensor - - import collections - from contextlib import contextmanager - from typing import Union, Optional, Dict, Tuple, Sequence - - _cache_enabled = 0 - _cache: Dict[Tuple[int, str], Optional[Tensor]] = {} - - @contextmanager - def cached(): - r"""Context manager that enables the caching system within parametrizations - registered with :func:`register_parametrization`. - - The value of the parametrized objects is computed and cached the first time - they are required when this context manager is active. The cached values are - discarded when leaving the context manager. - - This is useful when using a parametrized parameter more than once in the forward pass. - An example of this is when parametrizing the recurrent kernel of an RNN or when - sharing weights. - - The simplest way to activate the cache is by wrapping the forward pass of the neural network - - .. code-block:: python - - import torch.nn.utils.parametrize as P - ... - with P.cached(): - output = model(inputs) - - in training and evaluation. One may also wrap the parts of the modules that use - several times the parametrized tensors. For example, the loop of an RNN with a - parametrized recurrent kernel: - - .. code-block:: python - - with P.cached(): - for x in xs: - out_rnn = self.rnn_cell(x, out_rnn) - """ - global _cache - global _cache_enabled - _cache_enabled += 1 - try: - yield - finally: - _cache_enabled -= 1 - if not _cache_enabled: - _cache = {} - - def _register_parameter_or_buffer(module, name, X): - if isinstance(X, Parameter): - module.register_parameter(name, X) - else: - module.register_buffer(name, X) - - class ParametrizationList(ModuleList): - r"""A sequential container that holds and manages the ``original`` or ``original0``, ``original1``, ... - parameters or buffers of a parametrized :class:`torch.nn.Module`. - - It is the type of ``module.parametrizations[tensor_name]`` when ``module[tensor_name]`` - has been parametrized with :func:`register_parametrization`. - - If the first registered parmetrization has a ``right_inverse`` that returns one tensor or - does not have a ``right_inverse`` (in which case we assume that ``right_inverse`` is the identity), - it will hold the tensor under the name ``original``. - If it has a ``right_inverse`` that returns more than one tensor, these will be registered as - ``original0``, ``original1``, ... - - .. warning:: - This class is used internally by :func:`register_parametrization`. It is documented - here for completeness. It shall not be instantiated by the user. - - Args: - modules (sequence): sequence of modules representing the parametrizations - original (Parameter or Tensor): parameter or buffer that is parametrized - unsafe (bool): a boolean flag that denotes whether the parametrization - may change the dtype and shape of the tensor. Default: `False` - Warning: the parametrization is not checked for consistency upon registration. - Enable this flag at your own risk. - """ - - original: Tensor - unsafe: bool - - def __init__( - self, - modules: Sequence[Module], - original: Union[Tensor, Parameter], - unsafe: bool = False, - ) -> None: - # We require this because we need to treat differently the first parametrization - # This should never throw, unless this class is used from the outside - if len(modules) == 0: - raise ValueError("ParametrizationList requires one or more modules.") - - super().__init__(modules) - self.unsafe = unsafe - - # In plain words: - # module.weight must keep its dtype and shape. - # Furthermore, if there is no right_inverse or the right_inverse returns a tensor, - # this should be of the same dtype as the original tensor - # - # We check that the following invariants hold: - # X = module.weight - # Y = param.right_inverse(X) - # assert isinstance(Y, Tensor) or - # (isinstance(Y, collections.abc.Sequence) and all(isinstance(t, Tensor) for t in Y)) - # Z = param(Y) if isisntance(Y, Tensor) else param(*Y) - # # Consistency checks - # assert X.dtype == Z.dtype and X.shape == Z.shape - # # If it has one input, this allows to be able to use set_ to be able to - # # move data to/from the original tensor without changing its id (which is what the - # # optimiser uses to track parameters) - # if isinstance(Y, Tensor) - # assert X.dtype == Y.dtype - # Below we use original = X, new = Y - - original_shape = original.shape - original_dtype = original.dtype - - # Compute new - with torch.no_grad(): - new = original - for module in reversed(self): # type: ignore[call-overload] - if hasattr(module, "right_inverse"): - try: - new = module.right_inverse(new) - except NotImplementedError: - pass - # else, or if it throws, we assume that right_inverse is the identity - - if not isinstance(new, Tensor) and not isinstance( - new, collections.abc.Sequence - ): - raise ValueError( - "'right_inverse' must return a Tensor or a Sequence of tensors (list, tuple...). " - f"Got {type(new).__name__}" - ) - - # Set the number of original tensors - self.is_tensor = isinstance(new, Tensor) - self.ntensors = 1 if self.is_tensor else len(new) - - # Register the tensor(s) - if self.is_tensor: - if original.dtype != new.dtype: - raise ValueError( - "When `right_inverse` outputs one tensor, it may not change the dtype.\n" - f"original.dtype: {original.dtype}\n" - f"right_inverse(original).dtype: {new.dtype}" - ) - # Set the original to original so that the user does not need to re-register the parameter - # manually in the optimiser - with torch.no_grad(): - original.set_(new) # type: ignore[call-overload] - _register_parameter_or_buffer(self, "original", original) - else: - for i, originali in enumerate(new): - if not isinstance(originali, Tensor): - raise ValueError( - "'right_inverse' must return a Tensor or a Sequence of tensors " - "(list, tuple...). " - f"Got element {i} of the sequence with type {type(originali).__name__}." - ) - - # If the original tensor was a Parameter that required grad, we expect the user to - # add the new parameters to the optimizer after registering the parametrization - # (this is documented) - if isinstance(original, Parameter): - originali = Parameter(originali) - originali.requires_grad_(original.requires_grad) - _register_parameter_or_buffer(self, f"original{i}", originali) - - if not self.unsafe: - # Consistency checks: - # Since f : A -> B, right_inverse : B -> A, Z and original should live in B - # Z = forward(right_inverse(original)) - Z = self() - if not isinstance(Z, Tensor): - raise ValueError( - f"A parametrization must return a tensor. Got {type(Z).__name__}." - ) - if Z.dtype != original_dtype: - raise ValueError( - "Registering a parametrization may not change the dtype of the tensor, unless `unsafe` flag is enabled.\n" - f"unparametrized dtype: {original_dtype}\n" - f"parametrized dtype: {Z.dtype}" - ) - if Z.shape != original_shape: - raise ValueError( - "Registering a parametrization may not change the shape of the tensor, unless `unsafe` flag is enabled.\n" - f"unparametrized shape: {original_shape}\n" - f"parametrized shape: {Z.shape}" - ) - - def right_inverse(self, value: Tensor) -> None: - r"""Calls the methods ``right_inverse`` (see :func:`register_parametrization`) - of the parametrizations in the inverse order they were registered in. - Then, it stores the result in ``self.original`` if ``right_inverse`` outputs one tensor - or in ``self.original0``, ``self.original1``, ... if it outputs several. - - Args: - value (Tensor): Value to which initialize the module - """ - # All the exceptions in this function should almost never throw. - # They could throw if, for example, right_inverse function returns a different - # dtype when given a different input, which should most likely be caused by a - # bug in the user's code - - with torch.no_grad(): - # See https://github.com/pytorch/pytorch/issues/53103 - for module in reversed(self): # type: ignore[call-overload] - if hasattr(module, "right_inverse"): - value = module.right_inverse(value) - else: - raise RuntimeError( - f"parametrization {type(module).__name__} does not implement " - "right_inverse." - ) - if self.is_tensor: - # These exceptions should only throw when a right_inverse function does not - # return the same dtype for every input, which should most likely be caused by a bug - if not isinstance(value, Tensor): - raise ValueError( - f"`right_inverse` should return a tensor. Got {type(value).__name__}" - ) - if value.dtype != self.original.dtype: - raise ValueError( - f"The tensor returned by `right_inverse` has dtype {value.dtype} " - f"while `original` has dtype {self.original.dtype}" - ) - # We know that the result is going to have the same dtype - self.original.set_(value) # type: ignore[call-overload] - else: - if not isinstance(value, collections.abc.Sequence): - raise ValueError( - "'right_inverse' must return a sequence of tensors. " - f"Got {type(value).__name__}." - ) - if len(value) != self.ntensors: - raise ValueError( - "'right_inverse' must return a sequence of tensors of length " - f"{self.ntensors}. Got a sequence of lenght {len(value)}." - ) - for i, tensor in enumerate(value): - original_i = getattr(self, f"original{i}") - if not isinstance(tensor, Tensor): - raise ValueError( - f"`right_inverse` must return a sequence of tensors. " - f"Got element {i} of type {type(tensor).__name__}" - ) - if original_i.dtype != tensor.dtype: - raise ValueError( - f"Tensor {i} returned by `right_inverse` has dtype {tensor.dtype} " - f"while `original{i}` has dtype {original_i.dtype}" - ) - original_i.set_(tensor) - - def forward(self) -> Tensor: - # Unpack the originals for the first parametrization - if self.is_tensor: - x = self[0](self.original) - else: - originals = ( - getattr(self, f"original{i}") for i in range(self.ntensors) - ) - x = self[0](*originals) - # It's not possible to call self[1:] here, so we have to be a bit more cryptic - # Also we want to skip all non-integer keys - curr_idx = 1 - while hasattr(self, str(curr_idx)): - x = self[curr_idx](x) - curr_idx += 1 - return x - - def _inject_new_class(module: Module) -> None: - r"""Sets up a module to be parametrized. - - This works by substituting the class of the module by a class - that extends it to be able to inject a property - - Args: - module (nn.Module): module into which to inject the property - """ - cls = module.__class__ - - def getstate(self): - raise RuntimeError( - "Serialization of parametrized modules is only " - "supported through state_dict(). See:\n" - "https://pytorch.org/tutorials/beginner/saving_loading_models.html" - "#saving-loading-a-general-checkpoint-for-inference-and-or-resuming-training" - ) - - param_cls = type( - f"Parametrized{cls.__name__}", - (cls,), - { - "__getstate__": getstate, - }, - ) - - module.__class__ = param_cls - - def _inject_property(module: Module, tensor_name: str) -> None: - r"""Injects a property into module[tensor_name]. - - It assumes that the class in the module has already been modified from its - original one using _inject_new_class and that the tensor under :attr:`tensor_name` - has already been moved out - - Args: - module (nn.Module): module into which to inject the property - tensor_name (str): name of the name of the property to create - """ - # We check the precondition. - # This should never fire if register_parametrization is correctly implemented - assert not hasattr(module, tensor_name) - - @torch.jit.unused - def get_cached_parametrization(parametrization) -> Tensor: - global _cache - key = (id(module), tensor_name) - tensor = _cache.get(key) - if tensor is None: - tensor = parametrization() - _cache[key] = tensor - return tensor - - def get_parametrized(self) -> Tensor: - parametrization = self.parametrizations[tensor_name] - if _cache_enabled: - if torch.jit.is_scripting(): - # Scripting - raise RuntimeError( - "Caching is not implemented for scripting. " - "Either disable caching or avoid scripting." - ) - elif torch._C._get_tracing_state() is not None: - # Tracing - raise RuntimeError( - "Cannot trace a model while caching parametrizations." - ) - else: - return get_cached_parametrization(parametrization) - else: - # If caching is not active, this function just evaluates the parametrization - return parametrization() - - def set_original(self, value: Tensor) -> None: - self.parametrizations[tensor_name].right_inverse(value) - - setattr(module.__class__, tensor_name, property(get_parametrized, set_original)) - - def register_parametrization( - module: Module, - tensor_name: str, - parametrization: Module, - *, - unsafe: bool = False, - ) -> Module: - r"""Adds a parametrization to a tensor in a module. - - Assume that ``tensor_name="weight"`` for simplicity. When accessing ``module.weight``, - the module will return the parametrized version ``parametrization(module.weight)``. - If the original tensor requires a gradient, the backward pass will differentiate - through :attr:`parametrization`, and the optimizer will update the tensor accordingly. - - The first time that a module registers a parametrization, this function will add an attribute - ``parametrizations`` to the module of type :class:`~ParametrizationList`. - - The list of parametrizations on the tensor ``weight`` will be accessible under - ``module.parametrizations.weight``. - - The original tensor will be accessible under - ``module.parametrizations.weight.original``. - - Parametrizations may be concatenated by registering several parametrizations - on the same attribute. - - The training mode of a registered parametrization is updated on registration - to match the training mode of the host module - - Parametrized parameters and buffers have an inbuilt caching system that can be activated - using the context manager :func:`cached`. - - A :attr:`parametrization` may optionally implement a method with signature - - .. code-block:: python - - def right_inverse(self, X: Tensor) -> Union[Tensor, Sequence[Tensor]] - - This method is called on the unparametrized tensor when the first parametrization - is registered to compute the initial value of the original tensor. - If this method is not implemented, the original tensor will be just the unparametrized tensor. - - If all the parametrizations registered on a tensor implement `right_inverse` it is possible - to initialize a parametrized tensor by assigning to it, as shown in the example below. - - It is possible for the first parametrization to depend on several inputs. - This may be implemented returning a tuple of tensors from ``right_inverse`` - (see the example implementation of a ``RankOne`` parametrization below). - - In this case, the unconstrained tensors are also located under ``module.parametrizations.weight`` - with names ``original0``, ``original1``,... - - .. note:: - - If unsafe=False (default) both the forward and right_inverse methods will be called - once to perform a number of consistency checks. - If unsafe=True, then right_inverse will be called if the tensor is not parametrized, - and nothing will be called otherwise. - - .. note:: - - In most situations, ``right_inverse`` will be a function such that - ``forward(right_inverse(X)) == X`` (see - `right inverse `_). - Sometimes, when the parametrization is not surjective, it may be reasonable - to relax this. - - .. warning:: - - If a parametrization depends on several inputs, :func:`~register_parametrization` - will register a number of new parameters. If such parametrization is registered - after the optimizer is created, these new parameters will need to be added manually - to the optimizer. See :meth:`torch.Optimizer.add_param_group`. - - Args: - module (nn.Module): module on which to register the parametrization - tensor_name (str): name of the parameter or buffer on which to register - the parametrization - parametrization (nn.Module): the parametrization to register - Keyword args: - unsafe (bool): a boolean flag that denotes whether the parametrization - may change the dtype and shape of the tensor. Default: `False` - Warning: the parametrization is not checked for consistency upon registration. - Enable this flag at your own risk. - - Raises: - ValueError: if the module does not have a parameter or a buffer named :attr:`tensor_name` - - Examples: - >>> import torch - >>> import torch.nn as nn - >>> import torch.nn.utils.parametrize as P - >>> - >>> class Symmetric(nn.Module): - >>> def forward(self, X): - >>> return X.triu() + X.triu(1).T # Return a symmetric matrix - >>> - >>> def right_inverse(self, A): - >>> return A.triu() - >>> - >>> m = nn.Linear(5, 5) - >>> P.register_parametrization(m, "weight", Symmetric()) - >>> print(torch.allclose(m.weight, m.weight.T)) # m.weight is now symmetric - True - >>> A = torch.rand(5, 5) - >>> A = A + A.T # A is now symmetric - >>> m.weight = A # Initialize the weight to be the symmetric matrix A - >>> print(torch.allclose(m.weight, A)) - True - - >>> class RankOne(nn.Module): - >>> def forward(self, x, y): - >>> # Form a rank 1 matrix multiplying two vectors - >>> return x.unsqueeze(-1) @ y.unsqueeze(-2) - >>> - >>> def right_inverse(self, Z): - >>> # Project Z onto the rank 1 matrices - >>> U, S, Vh = torch.linalg.svd(Z, full_matrices=False) - >>> # Return rescaled singular vectors - >>> s0_sqrt = S[0].sqrt().unsqueeze(-1) - >>> return U[..., :, 0] * s0_sqrt, Vh[..., 0, :] * s0_sqrt - >>> - >>> linear_rank_one = P.register_parametrization(nn.Linear(4, 4), "weight", RankOne()) - >>> print(torch.linalg.matrix_rank(linear_rank_one.weight).item()) - 1 - - """ - parametrization.train(module.training) - if is_parametrized(module, tensor_name): - # Correctness checks. - # If A is the space of tensors with shape and dtype equal to module.weight - # we check that parametrization.forward and parametrization.right_inverse are - # functions from A to A - if not unsafe: - Y = getattr(module, tensor_name) - X = parametrization(Y) - if not isinstance(X, Tensor): - raise ValueError( - f"A parametrization must return a tensor. Got {type(X).__name__}." - ) - if X.dtype != Y.dtype: - raise ValueError( - "Registering a parametrization may not change the dtype of the tensor, unless the `unsafe` flag is enabled.\n" - f"module.{tensor_name}.dtype: {Y.dtype}\n" - f"parametrization(module.{tensor_name}).dtype: {X.dtype}" - ) - if X.shape != Y.shape: - raise ValueError( - "Registering a parametrization may not change the shape of the tensor, unless the `unsafe` flag is enabled.\n" - f"module.{tensor_name}.shape: {Y.shape}\n" - f"parametrization(module.{tensor_name}).shape: {X.shape}" - ) - if hasattr(parametrization, "right_inverse"): - try: - Z = parametrization.right_inverse(X) # type: ignore[operator] - except NotImplementedError: - pass - else: - if not isinstance(Z, Tensor): - raise ValueError( - f"parametrization.right_inverse must return a tensor. Got: {type(Z).__name__}" - ) - if Z.dtype != Y.dtype: - raise ValueError( - "The tensor returned by parametrization.right_inverse must have the same dtype " - f"as module.{tensor_name}, unless the `unsafe` flag is enabled.\n" - f"module.{tensor_name}.dtype: {Y.dtype}\n" - f"returned dtype: {Z.dtype}" - ) - if Z.shape != Y.shape: - raise ValueError( - "The tensor returned by parametrization.right_inverse must have the same shape " - f"as module.{tensor_name}, unless the `unsafe` flag is enabled.\n" - f"module.{tensor_name}.shape: {Y.shape}\n" - f"returned shape: {Z.shape}" - ) - # else right_inverse is assumed to be the identity - - # add the new parametrization to the parametrization list - assert isinstance(module.parametrizations, ModuleDict) # Make mypy happy - module.parametrizations[tensor_name].append(parametrization) - # If unsafe was True in previous parametrization, keep it enabled - module.parametrizations[tensor_name].unsafe |= unsafe # type: ignore[index, union-attr] - elif tensor_name in module._buffers or tensor_name in module._parameters: - # Set the parametrization mechanism - # Fetch the original buffer or parameter - original = getattr(module, tensor_name) - # We create this early to check for possible errors - parametrizations = ParametrizationList( - [parametrization], original, unsafe=unsafe - ) - # Delete the previous parameter or buffer - delattr(module, tensor_name) - # If this is the first parametrization registered on the module, - # we prepare the module to inject the property - if not is_parametrized(module): - # Change the class - _inject_new_class(module) - # Inject a ``ModuleDict`` into the instance under module.parametrizations - module.parametrizations = ModuleDict() - # Add a property into the class - _inject_property(module, tensor_name) - # Add a ParametrizationList - assert isinstance(module.parametrizations, ModuleDict) # Make mypy happy - module.parametrizations[tensor_name] = parametrizations - else: - raise ValueError( - f"Module '{module}' does not have a parameter, a buffer, or a " - f"parametrized element with name '{tensor_name}'" - ) - return module - - def is_parametrized(module: Module, tensor_name: Optional[str] = None) -> bool: - r"""Returns ``True`` if module has an active parametrization. - - If the argument :attr:`tensor_name` is specified, returns ``True`` if - ``module[tensor_name]`` is parametrized. - - Args: - module (nn.Module): module to query - name (str, optional): attribute in the module to query - Default: ``None`` - """ - parametrizations = getattr(module, "parametrizations", None) - if parametrizations is None or not isinstance(parametrizations, ModuleDict): - return False - if tensor_name is None: - # Check that there is at least one parametrized buffer or Parameter - return len(parametrizations) > 0 - else: - return tensor_name in parametrizations - - def remove_parametrizations( - module: Module, tensor_name: str, leave_parametrized: bool = True - ) -> Module: - r"""Removes the parametrizations on a tensor in a module. - - - If ``leave_parametrized=True``, ``module[tensor_name]`` will be set to - its current output. In this case, the parametrization shall not change the ``dtype`` - of the tensor. - - If ``leave_parametrized=False``, ``module[tensor_name]`` will be set to - the unparametrised tensor in ``module.parametrizations[tensor_name].original``. - This is only possible when the parametrization depends on just one tensor. - - Args: - module (nn.Module): module from which remove the parametrization - tensor_name (str): name of the parametrization to be removed - leave_parametrized (bool, optional): leave the attribute :attr:`tensor_name` parametrized. - Default: ``True`` - - Returns: - Module: module - - Raises: - ValueError: if ``module[tensor_name]`` is not parametrized - ValueError: if ``leave_parametrized=False`` and the parametrization depends on several tensors - """ - - if not is_parametrized(module, tensor_name): - raise ValueError( - f"Module {module} does not have a parametrization on {tensor_name}" - ) - - # Fetch the original tensor - assert isinstance(module.parametrizations, ModuleDict) # Make mypy happy - parametrizations = module.parametrizations[tensor_name] - if parametrizations.is_tensor: - original = parametrizations.original - if leave_parametrized: - with torch.no_grad(): - t = getattr(module, tensor_name) - # We know they have the same dtype because we have checked this when registering the - # parametrizations. As such, we can use set_ - # We do this so that the parameter does not to change the id() - # This way the user does not need to update the optimizer - with torch.no_grad(): - original.set_(t) - else: - if leave_parametrized: - # We cannot use no_grad because we need to know whether one or more - # original tensors required grad - t = getattr(module, tensor_name) - # We'll have to trust the user to add it to the optimizer - original = Parameter(t) if t.requires_grad else t - else: - raise ValueError( - "Cannot leave unparametrized (`leave_parametrized=False`) a tensor " - "that is parametrized in terms of a sequence of tensors." - ) - - # Delete the property that manages the parametrization - delattr(module.__class__, tensor_name) - # Delete the ParametrizationList - del module.parametrizations[tensor_name] - - # Restore the parameter / buffer into the main class - _register_parameter_or_buffer(module, tensor_name, original) - - # Roll back the parametrized class if no other buffer or parameter - # is currently parametrized in this class - if not is_parametrized(module): - delattr(module, "parametrizations") - # Restore class - orig_cls = module.__class__.__bases__[0] - module.__class__ = orig_cls - return module diff --git a/geotorch/pssdfixedrank.py b/geotorch/pssdfixedrank.py index e6dac445..62211b13 100644 --- a/geotorch/pssdfixedrank.py +++ b/geotorch/pssdfixedrank.py @@ -80,7 +80,7 @@ def sample(self, init_=torch.nn.init.xavier_normal_, eps=5e-6): >>> layer = nn.Linear(20, 20) >>> M = PSSD(layer.weight.size()) - >>> geotorch.register_parametrization(layer, "weight", M) + >>> torch.nn.utils.parametrize.register_parametrization(layer, "weight", M) >>> layer.weight = M.sample() Args: diff --git a/geotorch/skew.py b/geotorch/skew.py index c99c28ed..9fb080b2 100644 --- a/geotorch/skew.py +++ b/geotorch/skew.py @@ -19,6 +19,7 @@ def __init__(self, size, lower=True): self.n = n self.tensorial_size = tensorial_size self.lower = lower + self.register_buffer("_reference", torch.empty(0), persistent=False) @classmethod def parse_size(cls, size): @@ -79,7 +80,7 @@ def sample(self, init_=nn.init.xavier_normal_, lower=True): >>> layer = nn.Linear(20, 20) >>> M = Skew(layer.weight.size()) - >>> geotorch.register_parametrization(layer, "weight", M) + >>> torch.nn.utils.parametrize.register_parametrization(layer, "weight", M) >>> layer.weight = M.sample() Args: @@ -90,7 +91,7 @@ def sample(self, init_=nn.init.xavier_normal_, lower=True): Default: ``torch.nn.init.xavier_normal_`` """ with torch.no_grad(): - X = torch.empty(*(self.tensorial_size + (self.n, self.n))) + X = self._reference.new_empty(self.tensorial_size + (self.n, self.n)) init_(X) if lower: X.tril_(-1) diff --git a/geotorch/sl.py b/geotorch/sl.py index b9532a8a..173d19c8 100644 --- a/geotorch/sl.py +++ b/geotorch/sl.py @@ -1,6 +1,7 @@ import torch from .glp import GLp from .fixedrank import FixedRank +from .exceptions import InManifoldError class SL(GLp): @@ -35,18 +36,25 @@ def parse_f(f_name): def f_sl(x): y = f(x) - return y / y.prod(dim=-1, keepdim=True).pow(1.0 / y.shape[-1]) + log_y = y.log() + return (log_y - log_y.mean(dim=-1, keepdim=True)).exp() return (f_sl, inv) else: return f_name def in_manifold_singular_values(self, S, eps=5e-3): - if not super().in_manifold_singular_values(S, eps): + rank_eps = torch.finfo(S.dtype).eps * max(self.n, self.k) + if not super().in_manifold_singular_values(S, rank_eps): return False - # We compute the \infty-norm of the determinant minus 1 and should be about zero - infty_norm = (S.prod(dim=-1) - 1).abs().max(dim=-1).values - return (infty_norm < eps).all().item() + eps = max(eps, 8 * rank_eps**0.5) + logabsdet = S.log().sum(dim=-1).abs() + return (logabsdet < eps).all().item() + + def submersion_inv(self, X, check_in_manifold=True): + if check_in_manifold and not self.in_manifold(X): + raise InManifoldError(X, self) + return super().submersion_inv(X, check_in_manifold=False) def in_manifold(self, X, eps=5e-3): r""" @@ -58,8 +66,14 @@ def in_manifold(self, X, eps=5e-3): considered to be zero Default: ``5e-3`` """ - # The purpose of this function is just to have a more lax default eps value - return super().in_manifold(X, eps) + if X.size() != self.tensorial_size + (self.n, self.k): + return False + sign, logabsdet = torch.linalg.slogdet(X) + eps = max( + eps, + 8 * (torch.finfo(X.dtype).eps * max(self.n, self.k)) ** 0.5, + ) + return ((sign > 0) & (logabsdet.abs() < eps)).all().item() def sample(self, init_=torch.nn.init.xavier_normal_, eps=5e-6, factorized=False): r""" @@ -71,7 +85,7 @@ def sample(self, init_=torch.nn.init.xavier_normal_, eps=5e-6, factorized=False) >>> layer = nn.Linear(20, 20) >>> M = SL(layer.weight.size(), rank=6) - >>> geotorch.register_parametrization(layer, "weight", M) + >>> torch.nn.utils.parametrize.register_parametrization(layer, "weight", M) >>> layer.weight = M.sample() Args: @@ -82,8 +96,17 @@ def sample(self, init_=torch.nn.init.xavier_normal_, eps=5e-6, factorized=False) eps (float): Optional. Minimum singular value of the sampled matrix. Default: ``5e-6`` """ - U, S, V = super().sample(factorized=True, init_=init_) + U, S, V = super().sample(factorized=True, init_=init_, eps=eps) with torch.no_grad(): - # S >= 0, as given by torch.linalg.eigvalsh() - S = S / S.prod(dim=-1, keepdim=True).pow(1.0 / S.shape[-1]) - return (U * S.unsqueeze(-2)) @ V.transpose(-2, -1) + min_singular_value = max( + eps, + (torch.finfo(S.dtype).eps * max(self.n, self.k)) ** 0.5, + ) + S.clamp_min_(min_singular_value) + log_S = S.log() + S = (log_S - log_S.mean(dim=-1, keepdim=True)).exp() + X = (U * S.unsqueeze(-2)) @ V.transpose(-2, -1) + sign, logabsdet = torch.linalg.slogdet(X) + X[..., :, 0] *= sign.unsqueeze(-1) + X *= (-logabsdet / self.n).exp().unsqueeze(-1).unsqueeze(-1) + return X diff --git a/geotorch/so.py b/geotorch/so.py index 93b14ba7..bd494561 100644 --- a/geotorch/so.py +++ b/geotorch/so.py @@ -2,10 +2,7 @@ import torch from torch import nn -try: - from torch.linalg import matrix_exp as expm -except ImportError: - from torch import matrix_exp as expm +from torch.linalg import matrix_exp as expm from .utils import _extra_repr from .skew import Skew @@ -232,7 +229,7 @@ def torus_init_(tensor, init_=None, triv=expm): # Non-zero elements that we are going to set on the diagonal n_diag = n // 2 - diag = tensor.new(tensorial_size + (n_diag,)) + diag = tensor.new_empty(tensorial_size + (n_diag,)) if init_ is None: torch.nn.init.uniform_(diag, -math.pi, math.pi) else: diff --git a/geotorch/sphere.py b/geotorch/sphere.py index acb5ffd6..f14dcdb4 100644 --- a/geotorch/sphere.py +++ b/geotorch/sphere.py @@ -6,7 +6,7 @@ def project(x): - return x / x.norm(dim=-1, keepdim=True) + return x / torch.linalg.vector_norm(x, dim=-1, keepdim=True) def uniform_init_sphere_(x, r=1.0): @@ -16,34 +16,14 @@ def uniform_init_sphere_(x, r=1.0): """ with torch.no_grad(): x.normal_() - x.data = r * project(x) + x.copy_(r * project(x)) return x def _in_sphere(x, r, eps): - norm = x.norm(dim=-1) + norm = torch.linalg.vector_norm(x, dim=-1) rs = torch.full_like(norm, r) - return (torch.norm(norm - rs, p=float("inf")) < eps).all() - - -class sinc_class(torch.autograd.Function): - @staticmethod - def forward(ctx, x): - ctx.save_for_backward(x) - # Hardocoded for float, will do for now - ret = torch.sin(x) / x - ret[x.abs() < 1e-45] = 1.0 - return ret - - @staticmethod - def backward(ctx, grad_output): - (x,) = ctx.saved_tensors - ret = torch.cos(x) / x - torch.sin(x) / (x * x) - ret[x.abs() < 1e-10] = 0.0 - return ret * grad_output - - -sinc = sinc_class.apply + return (torch.linalg.vector_norm(norm - rs, ord=float("inf")) < eps).all() class SphereEmbedded(nn.Module): @@ -62,6 +42,7 @@ def __init__(self, size, radius=1.0): self.n = size[-1] self.tensorial_size = size[:-1] self.radius = SphereEmbedded.parse_radius(radius) + self.register_buffer("_reference", torch.empty(0), persistent=False) @staticmethod def parse_radius(radius): @@ -97,7 +78,7 @@ def sample(self): r""" Returns a uniformly sampled vector on the sphere. """ - x = torch.empty(*(self.tensorial_size) + (self.n,)) + x = self._reference.new_empty(self.tensorial_size + (self.n,)) return uniform_init_sphere_(x, r=self.radius) def extra_repr(self): @@ -132,7 +113,7 @@ def parse_radius(radius): return radius def frame(self, x, v): - projection = (v.unsqueeze(-2) @ x.unsqueeze(-1)).squeeze(-1) + projection = (v * x).sum(dim=-1, keepdim=True) v = v - projection * x return v @@ -140,8 +121,10 @@ def forward(self, v): x = self.base # Project v onto { = 0} v = self.frame(x, v) - vnorm = v.norm(dim=-1, keepdim=True) - return self.radius * (torch.cos(vnorm) * x + sinc(vnorm) * v) + vnorm = torch.linalg.vector_norm(v, dim=-1, keepdim=True) + return self.radius * ( + torch.cos(vnorm) * x + torch.special.sinc(vnorm / torch.pi) * v + ) def right_inverse(self, x, check_in_manifold=True): if check_in_manifold and not self.in_manifold(x): diff --git a/geotorch/symmetric.py b/geotorch/symmetric.py index b873ee93..f176e7b2 100644 --- a/geotorch/symmetric.py +++ b/geotorch/symmetric.py @@ -204,7 +204,7 @@ def sample(self, init_=torch.nn.init.xavier_normal_, factorized=False): >>> layer = nn.Linear(20, 20) >>> M = PSSD(layer.weight.size()) - >>> geotorch.register_parametrization(layer, "weight", M) + >>> torch.nn.utils.parametrize.register_parametrization(layer, "weight", M) >>> layer.weight = M.sample() Args: diff --git a/geotorch/utils.py b/geotorch/utils.py index c8d003be..ebe9f5f9 100644 --- a/geotorch/utils.py +++ b/geotorch/utils.py @@ -3,7 +3,7 @@ def update_base(layer, tensor_name): with torch.no_grad(): - setattr(layer, tensor_name, getattr(layer, tensor_name).data) + setattr(layer, tensor_name, getattr(layer, tensor_name).detach()) def transpose(fun): diff --git a/setup.cfg b/setup.cfg index eef7e9a0..dfa3fe48 100644 --- a/setup.cfg +++ b/setup.cfg @@ -3,11 +3,8 @@ # W503: black and flake8 disagree on how to place operators ignore = E203, W503 max-line-length = 88 -# Exclude as it will come directly from PyTorch in the future -exclude = geotorch/parametrize.py [coverage:report] omit = test/* setup.py - geotorch/parametrize.py diff --git a/test/test_hurwitz.py b/test/test_hurwitz.py index c9c05f3e..078571eb 100644 --- a/test/test_hurwitz.py +++ b/test/test_hurwitz.py @@ -1,18 +1,16 @@ import geotorch from geotorch.exceptions import NonSquareError, VectorError, InManifoldError from geotorch.hurwitz import Hurwitz, get_lyap_exp -from unittest import TestCase +from unittest import TestCase, skipUnless import torch -torch.set_default_dtype(torch.float64) - class TestHurwitz(TestCase): def setUp(self): self.size = (2, 5, 5) self.alpha = 1e-3 - self.hurwitz = Hurwitz(self.size, alpha=self.alpha) + self.hurwitz = Hurwitz(self.size, alpha=self.alpha).double() def test_parse_size_correct(self): n, tensorial_size = self.hurwitz.parse_size(self.size) @@ -33,34 +31,34 @@ def test_sample_shape_and_hurwitz(self): self.assertTrue(self.hurwitz.in_manifold_eigen(A)) def test_forward_and_submersion(self): - X1, X2, X3 = [torch.randn(5, 5) for _ in range(3)] + X1, X2, X3 = [torch.randn(*self.size, dtype=torch.float64) for _ in range(3)] A = self.hurwitz.forward(X1, X2, X3) self.assertEqual(A.shape, self.size) self.assertTrue(torch.is_tensor(A)) def test_submersion_inv_valid_matrix(self): A = self.hurwitz.sample() - alpha = float(get_lyap_exp(A) - 1e-4) - self.hurwitz.alpha = alpha - Q, P_inv, S = self.hurwitz.submersion_inv(A) + alpha = float(get_lyap_exp(A)) / 2 + inverse = Hurwitz(self.size, alpha=alpha).double() + Q, P_inv, S = inverse.submersion_inv(A) self.assertEqual(Q.shape, self.size) self.assertEqual(P_inv.shape, self.size) self.assertEqual(S.shape, self.size) - A_reconstructed = P_inv @ (-0.5 * Q + S) - self.hurwitz.alpha * torch.eye(5) + A_reconstructed = inverse.submersion(Q, P_inv, S) self.assertTrue(torch.allclose(A, A_reconstructed, atol=1e-4)) def test_submersion_inv_in_manifold_error(self): - A_unstable = torch.eye(5) + A_unstable = torch.eye(5, dtype=torch.float64).expand(self.size) with self.assertRaises(InManifoldError): self.hurwitz.submersion_inv(A_unstable) def test_right_inverse(self): A = self.hurwitz.sample() - self.hurwitz.alpha = float(get_lyap_exp(A)) - X1, X2, X3 = self.hurwitz.right_inverse(A) - A_from_ri = self.hurwitz.forward(X1, X2, X3) - print(torch.dist(A_from_ri, A)) + alpha = float(get_lyap_exp(A)) / 2 + inverse = Hurwitz(self.size, alpha=alpha).double() + X1, X2, X3 = inverse.right_inverse(A) + A_from_ri = inverse.forward(X1, X2, X3) self.assertTrue(torch.allclose(A, A_from_ri, atol=1e-4)) def test_in_manifold_eigen_true(self): @@ -73,7 +71,7 @@ def test_in_manifold_eigen_false(self): self.assertFalse(self.hurwitz.in_manifold_eigen(A)) def test_initialization_alpha_negative(self): - with self.assertRaises(AssertionError): + with self.assertRaises(ValueError): Hurwitz(self.size, alpha=-0.1) def test_extra_repr(self): @@ -88,19 +86,26 @@ def test_sample(self): self.assertTrue(hurwitz.in_manifold_eigen(A)) def test_batch_sample(self): - hurwitz = Hurwitz((4, 5, 5), alpha=self.alpha) - A = hurwitz.sample() + shapes = [] + + def init_(tensor): + shapes.append(tensor.shape) + return torch.nn.init.normal_(tensor) + + hurwitz = Hurwitz((4, 5, 5), alpha=self.alpha).double() + A = hurwitz.sample(init_) self.assertEqual(A.shape, (4, 5, 5)) + self.assertEqual(shapes, [torch.Size((4, 5, 5))] * 3) + self.assertFalse(torch.equal(A[0], A[1])) self.assertTrue(hurwitz.in_manifold_eigen(A)) def test_batch_submersion_inv(self): - hurwitz_batch = Hurwitz((7, 2, 2)) + hurwitz_batch = Hurwitz((7, 2, 2)).double() A_batch = hurwitz_batch.sample() - hurwitz_batch.alpha = float(get_lyap_exp(A_batch)) - Q, P_inv, S = hurwitz_batch.submersion_inv(A_batch) - A_reconstructed = P_inv @ (-0.5 * Q + S) - ( - hurwitz_batch.alpha - 1e-6 - ) * torch.eye(2) + alpha = float(get_lyap_exp(A_batch)) / 2 + inverse = Hurwitz((7, 2, 2), alpha=alpha).double() + Q, P_inv, S = inverse.submersion_inv(A_batch) + A_reconstructed = inverse.submersion(Q, P_inv, S) self.assertTrue(torch.allclose(A_batch, A_reconstructed, atol=1e-4)) @@ -110,3 +115,33 @@ def test_register_hurwitz(self): self.assertTrue( torch.all(torch.real(torch.linalg.eigvals(layer.weight)) <= -0.5) ) + + def test_sample_uses_manifold_dtype(self): + dtype = ( + torch.float64 + if torch.get_default_dtype() == torch.float32 + else torch.float32 + ) + hurwitz = Hurwitz((3, 3), alpha=self.alpha).to(dtype=dtype) + self.assertEqual(hurwitz.sample().dtype, dtype) + + def test_alpha_round_trips_in_state_dict(self): + source = Hurwitz((3, 3), alpha=0.75).double() + target = Hurwitz((3, 3), alpha=0.25).double() + target.load_state_dict(source.state_dict()) + self.assertEqual(target.alpha, source.alpha) + self.assertEqual(target.sample().dtype, source.sample().dtype) + + @skipUnless(torch.cuda.is_available(), "CUDA is not available") + def test_cuda_sampling_and_registration(self): + hurwitz = Hurwitz((2, 4, 4), alpha=0.5).to(device="cuda", dtype=torch.float64) + sample = hurwitz.sample() + self.assertEqual(sample.device.type, "cuda") + self.assertEqual(sample.dtype, torch.float64) + self.assertTrue(hurwitz.in_manifold_eigen(sample)) + + layer = torch.nn.Linear(4, 4, device="cuda", dtype=torch.float64) + geotorch.alpha_stable(layer, "weight", alpha=0.5) + self.assertEqual(layer.weight.device.type, "cuda") + self.assertEqual(layer.weight.dtype, torch.float64) + self.assertTrue(torch.all(torch.linalg.eigvals(layer.weight).real <= -0.5)) diff --git a/test/test_integration.py b/test/test_integration.py index 1943289a..c8db3404 100644 --- a/test/test_integration.py +++ b/test/test_integration.py @@ -5,7 +5,7 @@ import torch import torch.nn as nn -import geotorch.parametrize as P +from torch.nn.utils import parametrize import geotorch from geotorch.skew import Skew @@ -190,7 +190,7 @@ def _test_manifold(self, M, args_sample, args_constr, device, size, initialize): X = M_.sample(**args_sample) with torch.no_grad(): layer.weight.copy_(X) - P.register_parametrization(layer, "weight", M_) + parametrize.register_parametrization(layer, "weight", M_) # Check that it does not change the size of the layer self.assertEqual(old_size, layer.weight.size(), msg=f"{layer}") self._test_training(layer, args_sample, input_, initialize) @@ -203,7 +203,7 @@ def _test_training(self, layer, args_sample, input_, initialize): X = M.sample(**args_sample) self.assertTrue(M.in_manifold(X), msg=msg) layer.weight = X - with P.cached(): + with parametrize.cached(): # Compute the product if it is factorized # The sampled matrix should not have a gradient self.assertFalse(X.requires_grad) @@ -215,7 +215,7 @@ def _test_training(self, layer, args_sample, input_, initialize): # Take a couple SGD steps optim = torch.optim.SGD(layer.parameters(), lr=1e-3) for i in range(3): - with P.cached(): + with parametrize.cached(): loss = layer(input_).mean() optim.zero_grad() loss.backward() @@ -223,7 +223,7 @@ def _test_training(self, layer, args_sample, input_, initialize): # The layer stays in the manifold while being optimised self.assertTrue(M.in_manifold(layer.weight), msg=f"i:{i}\n" + msg) - with P.cached(): + with parametrize.cached(): weight_old = layer.weight update_base(layer, "weight") # After changing the base, the weight stays the same diff --git a/test/test_lowrank.py b/test/test_lowrank.py index f69d4a55..fd0b8e59 100644 --- a/test/test_lowrank.py +++ b/test/test_lowrank.py @@ -1,7 +1,9 @@ from unittest import TestCase +import torch + from geotorch.lowrank import LowRank -from geotorch.fixedrank import FixedRank +from geotorch.fixedrank import FixedRank, inv_softplus_epsilon, softplus_epsilon class TestLowRank(TestCase): @@ -22,3 +24,14 @@ def test_lowrank_errors(self): # On the wrong string with self.assertRaises(ValueError): FixedRank(size=(5, 3), rank=2, f="wrong") + + def test_fixed_rank_rejects_rank_deficient_matrix(self): + manifold = FixedRank(size=(3, 3), rank=3) + matrix = torch.diag(torch.tensor([3.0, 2.0, 0.0])) + self.assertFalse(manifold.in_manifold(matrix)) + + def test_inverse_softplus_is_stable(self): + x = torch.tensor([1e-4, 1.0, 100.0], dtype=torch.float64) + inverse = inv_softplus_epsilon(x) + self.assertTrue(torch.isfinite(inverse).all()) + self.assertTrue(torch.allclose(softplus_epsilon(inverse), x)) diff --git a/test/test_skew.py b/test/test_skew.py index 43225da9..ccd39bf3 100644 --- a/test/test_skew.py +++ b/test/test_skew.py @@ -4,8 +4,8 @@ import torch import torch.nn as nn +from torch.nn.utils import parametrize -import geotorch.parametrize as P from geotorch.skew import Skew import geotorch @@ -26,7 +26,7 @@ def test_backprop(self): # Assert that is stays in Skew(n) after some optimiser steps for _ in range(2): - with P.cached(): + with parametrize.cached(): self.assertTrue(Skew.in_manifold(layer.weight)) loss = layer(input_).sum() optim.zero_grad() @@ -47,3 +47,9 @@ def test_non_square(self): def test_repr(self): print(Skew((5, 5))) + + def test_sample_uses_module_dtype(self): + manifold = Skew((2, 3, 3)).double() + sample = manifold.sample() + self.assertEqual(sample.dtype, torch.float64) + self.assertEqual(sample.shape, (2, 3, 3)) diff --git a/test/test_sl.py b/test/test_sl.py index 58f0d0ba..bff649e0 100644 --- a/test/test_sl.py +++ b/test/test_sl.py @@ -1,5 +1,7 @@ from unittest import TestCase +import torch + from geotorch.sl import SL @@ -11,3 +13,26 @@ def test_SL_errors(self): # Try to instantiate it in a vector rather than a matrix with self.assertRaises(ValueError): SL(size=(5,)) + + def test_singular_values_are_normalized_in_log_domain(self): + manifold = SL(size=(3, 3)) + singular_values = manifold.f(torch.tensor([[-100.0, 0.0, 100.0]])) + self.assertTrue(torch.isfinite(singular_values).all()) + self.assertTrue( + torch.allclose( + singular_values.log().sum(dim=-1), + torch.zeros(1), + atol=1e-5, + ) + ) + + def test_negative_determinant_is_rejected(self): + manifold = SL(size=(2, 2)) + self.assertFalse(manifold.in_manifold(torch.diag(torch.tensor([-1.0, 1.0])))) + + def test_sample_has_positive_unit_determinant(self): + manifold = SL(size=(4, 5, 5)).double() + sample = manifold.sample() + self.assertTrue(manifold.in_manifold(sample)) + determinant = torch.linalg.det(sample) + self.assertTrue(torch.allclose(determinant, torch.ones_like(determinant))) diff --git a/test/test_sphere.py b/test/test_sphere.py index e9ba6839..5201fb7d 100644 --- a/test/test_sphere.py +++ b/test/test_sphere.py @@ -1,6 +1,9 @@ # Tests for the Sphere from unittest import TestCase +import pytest +import torch + from geotorch.sphere import Sphere, SphereEmbedded @@ -15,3 +18,36 @@ def test_construction(self): def test_repr(self): print(SphereEmbedded(size=(3,))) print(Sphere(size=(3,))) + + def test_autograd(self): + sphere = Sphere(size=(3,)).double() + v = (0.1 * torch.randn(3, dtype=torch.float64)).requires_grad_() + self.assertTrue(torch.autograd.gradcheck(sphere, (v,))) + + def test_torch_func(self): + sphere = Sphere(size=(3,)).double() + v = torch.zeros(3, dtype=torch.float64) + tangent = torch.randn_like(v) + + try: + torch.func.jacrev(lambda x: x.square())(v) + except NotImplementedError as error: + if "TensorWrapper" in str(error): + pytest.skip("torch.func is broken in this PyTorch build") + raise + + jacobian = torch.func.jacrev(sphere)(v) + output, jvp = torch.func.jvp(sphere, (v,), (tangent,)) + batched = torch.vmap(sphere)(torch.stack((v, tangent))) + + self.assertTrue(torch.isfinite(jacobian).all()) + self.assertTrue(torch.isfinite(output).all()) + self.assertTrue(torch.isfinite(jvp).all()) + self.assertTrue(torch.allclose(jvp, jacobian @ tangent)) + self.assertEqual(batched.shape, (2, 3)) + + def test_embedded_sample_uses_module_dtype(self): + sphere = SphereEmbedded(size=(2, 3)).double() + sample = sphere.sample() + self.assertEqual(sample.dtype, torch.float64) + self.assertEqual(sample.shape, (2, 3)) diff --git a/test/test_symmetric.py b/test/test_symmetric.py index c9a31e97..d9f5ac46 100644 --- a/test/test_symmetric.py +++ b/test/test_symmetric.py @@ -4,8 +4,8 @@ import torch import torch.nn as nn +from torch.nn.utils import parametrize -import geotorch.parametrize as P from geotorch.symmetric import Symmetric, SymF @@ -19,14 +19,16 @@ def test_backprop(self): for n, lower in itertools.product(sizes, [True, False]): layer = nn.Linear(n, n) - P.register_parametrization(layer, "weight", Symmetric(lower=lower)) + parametrize.register_parametrization( + layer, "weight", Symmetric(lower=lower) + ) input_ = torch.rand(5, n) optim = torch.optim.SGD(layer.parameters(), lr=1.0) # Assert that is stays in Sym(n) after some optimiser steps for _ in range(2): - with P.cached(): + with parametrize.cached(): self.assertTrue(Symmetric.in_manifold(layer.weight)) loss = layer(input_).sum() optim.zero_grad() From 7e4884a0b4c424d0a332053be3ffe497c7c43315 Mon Sep 17 00:00:00 2001 From: lezcano Date: Sat, 6 Jun 2026 21:16:24 +0100 Subject: [PATCH 2/6] Modernize package metadata and builds --- .readthedocs.yml | 9 ++++++- MANIFEST.in | 2 ++ README.rst | 5 ++-- geotorch/__init__.py | 2 +- pyproject.toml | 63 ++++++++++++++++++++++++++++++++++++++++++++ setup.cfg | 1 - setup.py | 55 -------------------------------------- 7 files changed, 77 insertions(+), 60 deletions(-) create mode 100644 MANIFEST.in create mode 100644 pyproject.toml delete mode 100644 setup.py diff --git a/.readthedocs.yml b/.readthedocs.yml index e7232557..7980310b 100644 --- a/.readthedocs.yml +++ b/.readthedocs.yml @@ -1,7 +1,14 @@ version: 2 +build: + os: ubuntu-24.04 + tools: + python: "3.13" + +sphinx: + configuration: docs/source/conf.py + python: - version: 3.7 install: - method: pip path: . diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 00000000..0e2d2876 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,2 @@ +prune examples +prune test diff --git a/README.rst b/README.rst index 712e33cc..7f22ea39 100644 --- a/README.rst +++ b/README.rst @@ -136,9 +136,10 @@ You may try GeoTorch installing it as .. code:: bash - pip install git+https://github.com/Lezcano/geotorch/ + pip install geotorch -GeoTorch is tested in Linux, Mac, and Windows environments for Python >= 3.6 and supports PyTorch >= 1.9 +GeoTorch supports Python >= 3.10 and PyTorch >= 2.6. It is tested on Linux, +macOS, and Windows. Sharing Weights, Parametrizations, and Normalizing Flows -------------------------------------------------------- diff --git a/geotorch/__init__.py b/geotorch/__init__.py index 05d2dfb7..8eb6066a 100644 --- a/geotorch/__init__.py +++ b/geotorch/__init__.py @@ -35,7 +35,7 @@ from .hurwitz import Hurwitz from .utils import update_base -__version__ = "0.3.0" +__version__ = "0.4.0" __all__ = [ diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 00000000..5881d99d --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,63 @@ +[build-system] +requires = ["setuptools>=77"] +build-backend = "setuptools.build_meta" + +[project] +name = "geotorch" +dynamic = ["version"] +description = "Constrained optimization and manifold optimization in PyTorch" +readme = "README.rst" +requires-python = ">=3.10" +license = "MIT" +license-files = ["LICENSE"] +authors = [ + {name = "Mario Lezcano Casado", email = "lezcano-93@hotmail.com"}, +] +keywords = [ + "constrained optimization", + "manifold optimization", + "PyTorch", +] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "Operating System :: OS Independent", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Topic :: Scientific/Engineering", + "Topic :: Scientific/Engineering :: Mathematics", +] +dependencies = ["torch>=2.6"] + +[project.optional-dependencies] +test = [ + "pytest", + "pytest-xdist", +] +dev = [ + "black", + "flake8", + "pytest", + "pytest-cov", + "pytest-xdist", + "sphinx", + "sphinx-autodoc-typehints", + "sphinx-rtd-theme", + "sphinxcontrib-spelling", +] + +[project.urls] +Documentation = "https://geotorch.readthedocs.io/" +Repository = "https://github.com/Lezcano/geotorch" + +[tool.setuptools.dynamic] +version = {attr = "geotorch.__version__"} + +[tool.setuptools.packages.find] +include = ["geotorch*"] +exclude = ["examples*", "test*"] diff --git a/setup.cfg b/setup.cfg index dfa3fe48..e941eadb 100644 --- a/setup.cfg +++ b/setup.cfg @@ -7,4 +7,3 @@ max-line-length = 88 [coverage:report] omit = test/* - setup.py diff --git a/setup.py b/setup.py deleted file mode 100644 index 276df470..00000000 --- a/setup.py +++ /dev/null @@ -1,55 +0,0 @@ -from setuptools import setup, find_packages -import os -import re - -TEST_REQUIRES = ["pytest"] - -DEV_REQUIRES = TEST_REQUIRES + [ - "black", - "flake8", - "sphinx", - "sphinx-autodoc-typehints", - "sphinx-rtd-theme", - "sphinxcontrib-spelling", - "codecov", -] - -classifiers = [ - "Development Status :: 3 - Alpha", - "Programming Language :: Python", - "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.5", - "License :: OSI Approved :: MIT License", - "Topic :: Scientific/Engineering", - "Topic :: Scientific/Engineering :: Mathematics", - "Intended Audience :: Science/Research", - "Intended Audience :: Developers", - "Operating System :: OS Independent", -] - -# Get the long description from the README file -with open("README.rst", "r", encoding="utf8") as fh: - long_description = fh.read() - -# Get version string from module -init_path = os.path.join(os.path.dirname(__file__), "geotorch/__init__.py") -with open(init_path, "r", encoding="utf8") as f: - version = re.search(r"__version__ = ['\"]([^'\"]*)['\"]", f.read(), re.M).group(1) - -setup( - name="geotorch", - version=version, - description="Constrained Optimization and Manifold Optimization in Pytorch", - author="Mario Lezcano Casado", - author_email="lezcano-93@hotmail.com", - license="MIT", - long_description=long_description, - long_description_content_type="text/x-rst", - url="https://github.com/Lezcano/geotorch", - classifiers=classifiers, - keywords=["Constrained Optimization", "Optimization on Manifolds", "Pytorch"], - packages=find_packages(), - python_requires=">=3.5", - install_requires=["torch>=1.9"], - extras_require={"dev": DEV_REQUIRES, "test": TEST_REQUIRES}, -) From 5c671b03764a0edf20512ffdebab791ad978dca1 Mon Sep 17 00:00:00 2001 From: lezcano Date: Sat, 6 Jun 2026 21:16:36 +0100 Subject: [PATCH 3/6] Modernize continuous integration --- .github/scripts/check_dist.py | 55 ++++++++++++++ .github/workflows/build.yml | 128 +++++++++++++++++++++++---------- .github/workflows/coverage.yml | 35 --------- README.rst | 4 +- 4 files changed, 146 insertions(+), 76 deletions(-) create mode 100644 .github/scripts/check_dist.py delete mode 100644 .github/workflows/coverage.yml diff --git a/.github/scripts/check_dist.py b/.github/scripts/check_dist.py new file mode 100644 index 00000000..7067f01e --- /dev/null +++ b/.github/scripts/check_dist.py @@ -0,0 +1,55 @@ +import argparse +from email.parser import BytesParser +from pathlib import Path, PurePosixPath +import tarfile +import zipfile + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("dist", nargs="+", type=Path) + parser.add_argument("--tag") + return parser.parse_args() + + +def archive_names(path): + if path.suffix == ".whl": + with zipfile.ZipFile(path) as archive: + return archive.namelist() + with tarfile.open(path) as archive: + return archive.getnames() + + +def wheel_version(path): + with zipfile.ZipFile(path) as archive: + metadata = next( + name for name in archive.namelist() if name.endswith(".dist-info/METADATA") + ) + message = BytesParser().parsebytes(archive.read(metadata)) + return message["Version"] + + +def main(): + args = parse_args() + wheels = [path for path in args.dist if path.suffix == ".whl"] + sdists = [path for path in args.dist if path.name.endswith(".tar.gz")] + if len(wheels) != 1 or len(sdists) != 1: + raise SystemExit("Expected exactly one wheel and one sdist") + + for path in args.dist: + for name in archive_names(path): + parts = PurePosixPath(name).parts + if "examples" in parts or "test" in parts: + raise SystemExit(f"Unexpected path in {path.name}: {name}") + + version = wheel_version(wheels[0]) + if f"-{version}.tar.gz" not in sdists[0].name: + raise SystemExit("Wheel and sdist versions do not match") + if args.tag is not None and args.tag.removeprefix("v") != version: + raise SystemExit(f"Tag {args.tag!r} does not match package version {version!r}") + + print(f"Validated GeoTorch {version} wheel and sdist") + + +if __name__ == "__main__": + main() diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 7abf216d..92d71bd3 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -1,52 +1,102 @@ -name: build +name: CI on: push: - branches: [ master ] + branches: [master] pull_request: - branches: [ '*' ] + +concurrency: + group: ci-${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read jobs: - build: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: "3.14" + cache: pip + - run: python -m pip install --upgrade pip + - run: python -m pip install ".[dev]" + - run: black --check --diff . + - run: flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics + + linux: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + include: + - python: "3.10" + torch: "torch==2.6.*" + - python: "3.12" + torch: "torch" + - python: "3.14" + torch: "torch" + steps: + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: ${{ matrix.python }} + cache: pip + - run: python -m pip install --upgrade pip + - run: python -m pip install "${{ matrix.torch }}" pytest pytest-xdist + - run: python -m pip install --no-deps -e . + - run: python -m pytest -s --tb=short -n32 test + os-smoke: runs-on: ${{ matrix.os }} strategy: + fail-fast: false matrix: - python-version: [3.6, 3.7, 3.8] os: [ubuntu-latest, macos-latest, windows-latest] + steps: + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: "3.13" + cache: pip + - run: python -m pip install --upgrade pip + - run: python -m pip install ".[test]" + - run: python -m pytest -s --tb=short -n32 test + + coverage: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: "3.14" + cache: pip + - run: python -m pip install --upgrade pip + - run: python -m pip install ".[test]" pytest-cov + - run: python -m pytest -s --tb=short -n32 --cov=geotorch --cov-report=xml test + - uses: codecov/codecov-action@e79a6962e0d4c0c17b229090214935d2e33f8354 # v6.0.1 + with: + token: ${{ secrets.CODECOV_TOKEN }} + package: + runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 - with: - python-version: ${{ matrix.python-version }} - - - name: Install dependencies others - if: ${{ matrix.os != 'windows-latest' }} - run: | - python -m pip install --upgrade pip - pip install .[dev] - - # Windows is treated differently, as PyTorch is not uploaded to Pypi atm - - name: Install dependencies windows - if: ${{ matrix.os == 'windows-latest' }} - run: | - python -m pip install --upgrade pip - pip install torch===1.10.0+cpu -f https://download.pytorch.org/whl/torch_stable.html - pip install .[dev] - - - name: Lint with flake8 - run: | - # stop the build if there are Python syntax errors or undefined names - flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics - # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide - flake8 . --count --exit-zero --max-complexity=12 --max-line-length=127 --statistics - - - name: Lint with Black - run: | - black --check --diff . - - - name: Test with pytest - run: | - pytest test + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: "3.14" + cache: pip + - run: python -m pip install --upgrade pip + - run: python -m pip install build pytest pytest-xdist torch twine + - run: python -m build + - run: python -m twine check dist/* + - run: python .github/scripts/check_dist.py dist/* + - name: Test installed wheel + run: | + python -m venv --system-site-packages /tmp/geotorch-wheel + /tmp/geotorch-wheel/bin/python -m pip install --no-deps dist/*.whl + cp -R test /tmp/geotorch-tests + cd /tmp + /tmp/geotorch-wheel/bin/python -m pytest -s --tb=short -n32 geotorch-tests diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml deleted file mode 100644 index 3b94452a..00000000 --- a/.github/workflows/coverage.yml +++ /dev/null @@ -1,35 +0,0 @@ -name: Coverage - -on: - push: - branches: [ master ] - pull_request: - branches: [ master ] - -jobs: - build: - - runs-on: ubuntu-latest - - # Setup lastest python version - steps: - - uses: actions/checkout@v2 - - name: Set up Python 3.8 - uses: actions/setup-python@v2 - with: - python-version: '3.8' - - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install codecov pytest-cov - pip install .[dev] - - - name: Coverage tests - run: pytest --cov=geotorch test/ --cov-report term-missing - - - name: Codecov - if: success() - env: - CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} - run: codecov diff --git a/README.rst b/README.rst index 7f22ea39..44dc9c8c 100644 --- a/README.rst +++ b/README.rst @@ -211,8 +211,8 @@ Please cite the following work if you found GeoTorch useful. This paper exposes } -.. |Build| image:: https://github.com/lezcano/geotorch/workflows/Build/badge.svg - :target: https://github.com/lezcano/geotorch/workflows/Build/badge.svg +.. |Build| image:: https://github.com/lezcano/geotorch/actions/workflows/build.yml/badge.svg + :target: https://github.com/lezcano/geotorch/actions/workflows/build.yml :alt: Build .. |Docs| image:: https://readthedocs.org/projects/geotorch/badge/?version=latest :target: https://geotorch.readthedocs.io/en/latest/?badge=latest From 7e4ad9336aa7ab66236c2da6a10e35da75273b7d Mon Sep 17 00:00:00 2001 From: lezcano Date: Sat, 6 Jun 2026 21:16:45 +0100 Subject: [PATCH 4/6] Publish releases through PyPI trusted publishing --- .github/workflows/pypi.yml | 66 +++++++++++++++++++++++++++----------- 1 file changed, 47 insertions(+), 19 deletions(-) diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml index d90a01cc..203af691 100644 --- a/.github/workflows/pypi.yml +++ b/.github/workflows/pypi.yml @@ -1,26 +1,54 @@ -name: Upload Python Package +name: Publish to PyPI on: release: - types: [created] + types: [published] + +permissions: + contents: read jobs: - deploy: + build: + if: github.event.release.prerelease == false + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 + with: + fetch-depth: 0 + ref: ${{ github.event.release.tag_name }} + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: "3.14" + cache: pip + - name: Verify release source + env: + RELEASE_TAG: ${{ github.event.release.tag_name }} + run: | + test "$(git rev-parse HEAD)" = "$(git rev-list -n 1 "$RELEASE_TAG")" + - run: python -m pip install --upgrade pip + - run: python -m pip install build twine + - run: python -m build + - run: python -m twine check dist/* + - name: Verify release artifacts + env: + RELEASE_TAG: ${{ github.event.release.tag_name }} + run: python .github/scripts/check_dist.py --tag "$RELEASE_TAG" dist/* + - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: distributions + path: dist/* + if-no-files-found: error + + publish: + if: github.event.release.prerelease == false + needs: build runs-on: ubuntu-latest + environment: pypi + permissions: + id-token: write steps: - - uses: actions/checkout@v2 - - name: Set up Python 3.x - uses: actions/setup-python@v2 - with: - python-version: '3.x' - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install setuptools wheel twine - - name: Build and publish - env: - TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }} - TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }} - run: | - python setup.py sdist bdist_wheel - twine upload dist/* + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: distributions + path: dist + - uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0 From e5835f718f8fc7843521d0b2de051321250df27f Mon Sep 17 00:00:00 2001 From: lezcano Date: Sat, 6 Jun 2026 21:39:54 +0100 Subject: [PATCH 5/6] Use stable equivalent scalar expressions --- examples/sequential_mnist.py | 4 +--- geotorch/almostorthogonal.py | 4 ++-- test/test_almostorthogonal.py | 14 +++++++++++++- 3 files changed, 16 insertions(+), 6 deletions(-) diff --git a/examples/sequential_mnist.py b/examples/sequential_mnist.py index 77989654..b2510dc6 100644 --- a/examples/sequential_mnist.py +++ b/examples/sequential_mnist.py @@ -100,9 +100,7 @@ def reset_parameters(self): # seems to help in classification problems def init_(x): - x.uniform_(0.0, math.pi / 2.0) - c = torch.cos(x) - x.copy_(-torch.sqrt((1.0 - c) / (1.0 + c))) + x.uniform_(0.0, math.pi / 2.0).mul_(0.5).tan_().neg_() K = self.recurrent_kernel # We initialize it by assigning directly to it from a sampler diff --git a/geotorch/almostorthogonal.py b/geotorch/almostorthogonal.py index f38720b5..56626d29 100644 --- a/geotorch/almostorthogonal.py +++ b/geotorch/almostorthogonal.py @@ -5,11 +5,11 @@ def scaled_sigmoid(t): - return 2.0 * (torch.sigmoid(t) - 0.5) + return torch.tanh(0.5 * t) def inv_scaled_sigmoid(t): - return torch.logit(0.5 * t + 0.5) + return 2.0 * torch.atanh(t) class AlmostOrthogonal(LowRank): diff --git a/test/test_almostorthogonal.py b/test/test_almostorthogonal.py index 25004fb3..b6cd2e75 100644 --- a/test/test_almostorthogonal.py +++ b/test/test_almostorthogonal.py @@ -1,6 +1,12 @@ from unittest import TestCase -from geotorch.almostorthogonal import AlmostOrthogonal +import torch + +from geotorch.almostorthogonal import ( + AlmostOrthogonal, + inv_scaled_sigmoid, + scaled_sigmoid, +) class TestLowRank(TestCase): @@ -21,3 +27,9 @@ def test_almostorthogonal_errors(self): # Or too small with self.assertRaises(ValueError): AlmostOrthogonal(size=(5, 4), lam=-1.0) + + def test_scaled_sigmoid_is_accurate_near_zero(self): + x = torch.tensor([1e-8], dtype=torch.float32) + y = scaled_sigmoid(x) + self.assertNotEqual(y.item(), 0.0) + torch.testing.assert_close(inv_scaled_sigmoid(y), x) From 807781cbe9614db319c85e230b2df0ab738ff963 Mon Sep 17 00:00:00 2001 From: lezcano Date: Sat, 6 Jun 2026 21:43:01 +0100 Subject: [PATCH 6/6] Reduce tensor temporaries in manifold operations --- examples/copying_problem.py | 9 ++++----- examples/eigenvalue.py | 2 +- examples/sequential_mnist.py | 15 ++++++++------- geotorch/fixedrank.py | 5 +---- geotorch/grassmannian.py | 9 +++------ geotorch/lowrank.py | 14 +++++--------- geotorch/so.py | 6 +----- geotorch/sphere.py | 4 +--- geotorch/stiefel.py | 4 ++-- geotorch/symmetric.py | 12 ++++-------- 10 files changed, 30 insertions(+), 50 deletions(-) diff --git a/examples/copying_problem.py b/examples/copying_problem.py index 4fef12f8..316a4cf8 100644 --- a/examples/copying_problem.py +++ b/examples/copying_problem.py @@ -118,7 +118,7 @@ def forward(self, inputs): return torch.stack(outputs, dim=1) def loss(self, logits, y): - return self.loss_func(logits.view(-1, 9), y.view(-1)) + return self.loss_func(logits.flatten(0, 1), y.flatten()) def accuracy(self, logits, y): return torch.eq(torch.argmax(logits, dim=2), y).float().mean() @@ -155,10 +155,9 @@ def main(): model = Model(alphabet_size, hidden_size).to(device) p_orth = model.rnn.recurrent_kernel - orth_params = p_orth.parameters() - non_orth_params = ( - p for p in model.parameters() if p not in set(p_orth.parameters()) - ) + orth_params = list(p_orth.parameters()) + orth_params_set = set(orth_params) + non_orth_params = [p for p in model.parameters() if p not in orth_params_set] if RGD: # Implement Stochstic Riemannian Gradient Descent via SGD diff --git a/examples/eigenvalue.py b/examples/eigenvalue.py index 3cb6cc50..1245b95b 100644 --- a/examples/eigenvalue.py +++ b/examples/eigenvalue.py @@ -24,7 +24,7 @@ def __init__(self, n): def forward(self, A): x = self.x - return x @ A @ x + return torch.linalg.multi_dot((x, A, x)) # Generate matrix diff --git a/examples/sequential_mnist.py b/examples/sequential_mnist.py index b2510dc6..cd74ca39 100644 --- a/examples/sequential_mnist.py +++ b/examples/sequential_mnist.py @@ -139,7 +139,7 @@ def loss(self, logits, y): return self.loss_func(logits, y) def correct(self, logits, y): - return torch.eq(torch.argmax(logits, dim=1), y).float().sum() + return torch.eq(torch.argmax(logits, dim=1), y).sum() def main(): @@ -166,10 +166,11 @@ def main(): model.train() p_orth = model.rnn.recurrent_kernel - orth_params = p_orth.parameters() - non_orth_params = ( - param for param in model.parameters() if param not in set(p_orth.parameters()) - ) + orth_params = list(p_orth.parameters()) + orth_params_set = set(orth_params) + non_orth_params = [ + param for param in model.parameters() if param not in orth_params_set + ] optim = torch.optim.RMSprop( [{"params": non_orth_params}, {"params": orth_params, "lr": args.lr_orth}], @@ -180,7 +181,7 @@ def main(): for epoch in range(epochs): processed = 0 for batch_idx, (batch_x, batch_y) in enumerate(train_loader): - batch_x, batch_y = batch_x.to(device).view(-1, 784), batch_y.to(device) + batch_x, batch_y = batch_x.to(device).flatten(1), batch_y.to(device) optim.zero_grad() logits = model(batch_x) @@ -209,7 +210,7 @@ def main(): test_loss = 0.0 correct = 0.0 for batch_x, batch_y in test_loader: - batch_x, batch_y = batch_x.to(device).view(-1, 784), batch_y.to(device) + batch_x, batch_y = batch_x.to(device).flatten(1), batch_y.to(device) logits = model(batch_x) test_loss += model.loss(logits, batch_y).float() correct += model.correct(logits, batch_y).float() diff --git a/geotorch/fixedrank.py b/geotorch/fixedrank.py index 7b71d780..70d3e511 100644 --- a/geotorch/fixedrank.py +++ b/geotorch/fixedrank.py @@ -119,7 +119,4 @@ def sample(self, init_=torch.nn.init.xavier_normal_, eps=5e-6, factorized=False) return U, S, V else: # Compute U S V^T efficiently - if self.transposed: - return (U * S.unsqueeze(-2)) @ V.transpose(-2, -1) - else: - return U @ (S.unsqueeze(-1) * V.transpose(-2, -1)) + return (U * S.unsqueeze(-2)) @ V.mT diff --git a/geotorch/grassmannian.py b/geotorch/grassmannian.py index 995f6158..65b551c1 100644 --- a/geotorch/grassmannian.py +++ b/geotorch/grassmannian.py @@ -1,4 +1,4 @@ -import torch +from torch.nn import functional as F from .stiefel import Stiefel @@ -20,8 +20,5 @@ def __init__(self, size, triv="expm"): super().__init__(size=size, triv=triv) def frame(self, X): - k = X.size(-1) - size_z = X.size()[:-2] + (k, k) - Z = X.new_zeros(*size_z) - X = torch.cat([Z, X[..., k:, :]], dim=-2) - return super().frame(X) + n, k = X.size(-2), X.size(-1) + return F.pad(X[..., k:, :], (0, n - k, k, 0)) diff --git a/geotorch/lowrank.py b/geotorch/lowrank.py index 90c09cd0..b93ea932 100644 --- a/geotorch/lowrank.py +++ b/geotorch/lowrank.py @@ -71,10 +71,10 @@ def frame_inv(self, X1, X2, X3): # X2 is a vector # X3 is lower-triangular size = self.tensorial_size + (self.n, self.k) - ret = torch.zeros(size, dtype=X1.dtype, device=X1.device) - ret[..., : self.rank] += X1 - ret[..., : self.rank, : self.rank] += torch.diag_embed(X2) - ret.transpose(-2, -1)[..., : self.rank] += X3 + ret = X1.new_zeros(size) + ret[..., : self.rank].copy_(X1) + ret.diagonal(dim1=-2, dim2=-1)[..., : self.rank].add_(X2) + ret.mT[..., : self.rank].add_(X3) return ret def submersion_inv(self, X, check_in_manifold=True): @@ -147,11 +147,7 @@ def sample(self, init_=torch.nn.init.xavier_normal_, factorized=False): Default: ``torch.nn.init.xavier_normal_`` """ with torch.no_grad(): - device = self[0].base.device - dtype = self[0].base.dtype - X = torch.empty( - *(self.tensorial_size + (self.n, self.k)), device=device, dtype=dtype - ) + X = self[0].base.new_empty(self.tensorial_size + (self.n, self.k)) init_(X) U, S, Vt = torch.linalg.svd(X, full_matrices=False) U, S, Vt = U[..., : self.rank], S[..., : self.rank], Vt[..., : self.rank, :] diff --git a/geotorch/so.py b/geotorch/so.py index bd494561..c8836f6b 100644 --- a/geotorch/so.py +++ b/geotorch/so.py @@ -135,11 +135,7 @@ def sample(self, distribution="uniform", init_=None): `torch.init `_. Default: :math:`\operatorname{Uniform}(-\pi, \pi)` """ - device = self.base.device - dtype = self.base.dtype - ret = torch.empty( - *(self.tensorial_size + (self.n, self.n)), device=device, dtype=dtype - ) + ret = self.base.new_empty(self.tensorial_size + (self.n, self.n)) if distribution == "uniform": uniform_init_(ret) elif distribution == "torus": diff --git a/geotorch/sphere.py b/geotorch/sphere.py index f14dcdb4..71de4cb3 100644 --- a/geotorch/sphere.py +++ b/geotorch/sphere.py @@ -152,9 +152,7 @@ def sample(self): r""" Returns a uniformly sampled vector on the sphere. """ - device = self.base.device - dtype = self.base.dtype - x = torch.empty(*(self.tensorial_size) + (self.n,), device=device, dtype=dtype) + x = self.base.new_empty(self.tensorial_size + (self.n,)) return uniform_init_sphere_(x, r=self.radius) def extra_repr(self): diff --git a/geotorch/stiefel.py b/geotorch/stiefel.py index 3b8cb302..2b209740 100644 --- a/geotorch/stiefel.py +++ b/geotorch/stiefel.py @@ -1,4 +1,5 @@ import torch +from torch.nn import functional as F from .utils import transpose, _extra_repr from .so import SO, _has_orthonormal_columns @@ -34,8 +35,7 @@ def size_so(cls, size): def frame(self, X): n, k = X.size(-2), X.size(-1) - size_z = X.size()[:-2] + (n, n - k) - return torch.cat([X, X.new_zeros(*size_z)], dim=-1) + return F.pad(X, (0, n - k)) @transpose def forward(self, X): diff --git a/geotorch/symmetric.py b/geotorch/symmetric.py index f176e7b2..45d9487a 100644 --- a/geotorch/symmetric.py +++ b/geotorch/symmetric.py @@ -129,10 +129,10 @@ def forward(self, X): def frame_inv(self, X1, X2): size = self.tensorial_size + (self.n, self.n) - ret = torch.zeros(*size, dtype=X1.dtype, device=X1.device) + ret = X1.new_zeros(size) with torch.no_grad(): - ret[..., : self.rank] += X1 - ret[..., : self.rank, : self.rank] += torch.diag_embed(X2) + ret[..., : self.rank].copy_(X1) + ret.diagonal(dim1=-2, dim2=-1)[..., : self.rank].add_(X2) return ret def submersion_inv(self, X, check_in_manifold=True): @@ -215,11 +215,7 @@ def sample(self, init_=torch.nn.init.xavier_normal_, factorized=False): Default: ``torch.nn.init.xavier_normal_`` """ with torch.no_grad(): - device = self[0].base.device - dtype = self[0].base.dtype - X = torch.empty( - *(self.tensorial_size + (self.n, self.n)), device=device, dtype=dtype - ) + X = self[0].base.new_empty(self.tensorial_size + (self.n, self.n)) init_(X) X = X @ X.transpose(-2, -1) L, Q = torch.linalg.eigh(X)