From 6ffce16072b98b5102ab9bbfa429f15d9ed8e24e Mon Sep 17 00:00:00 2001 From: Ahmed Eldeeb <62363199+deeb01@users.noreply.github.com> Date: Mon, 14 Sep 2026 13:49:51 -0700 Subject: [PATCH 1/3] Add a PyTorch solver to ot.dr.wda (#806) Adds solver='torch' to wda: PyTorch autodiff with Riemannian gradient descent on the Stiefel manifold, using a QR retraction and backtracking with an adaptive initial step. It mirrors what pymanopt's SteepestDescent does, so both solvers target the same optimum rather than two different algorithms. The torch path needs only torch, so it works on installations without autograd or pymanopt, and accepts torch tensors directly, keeping their device and dtype. To make that possible, ot.dr's dependencies are now imported optionally and each function raises an ImportError naming what it needs, rather than the module failing to import unless all of them are present. Verified that the torch objective and its gradient match the autograd ones at the same point, and that both solvers reach a comparable objective from the same starting point. Also raises a clear ValueError when the between-class transport cost underflows to zero, which previously produced a divide-by-zero warning and an undefined objective. --- RELEASES.md | 7 ++ ot/dr.py | 249 ++++++++++++++++++++++++++++++++++++++++++++++-- test/test_dr.py | 147 +++++++++++++++++++++++++++- 3 files changed, 394 insertions(+), 9 deletions(-) diff --git a/RELEASES.md b/RELEASES.md index d636e8cf8..657bf11de 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -4,6 +4,13 @@ #### New features +- `ot.dr.wda` gains `solver='torch'`, a PyTorch autodiff solver with Riemannian gradient descent, usable on installations without autograd or pymanopt, and accepting torch tensors directly (PR #853, Issue #806) +- `ot.dr` dependencies (autograd, pymanopt, scikit-learn, torch) are now imported optionally, so importing `ot.dr` no longer requires all of them; each function raises an explicit `ImportError` naming what it needs (PR #853) + +## 0.9.8dev + +#### New features + - Add stereographic spherical sliced Wasserstein distance in `ot.sliced.stereographic_sliced_wasserstein_sphere`, with its rotationally invariant extension (PR #836) - Add Quasi-Monte Carlo sliced Wasserstein sampling (QSW/RQSW) via generalized spiral points, selectable with `sampling_slices` in `sliced_wasserstein_distance`, diff --git a/ot/dr.py b/ot/dr.py index 9cfabff5f..a3922ee22 100644 --- a/ot/dr.py +++ b/ot/dr.py @@ -18,22 +18,54 @@ from scipy import linalg +# ot.dr offers solvers with different dependencies. Each is imported optionally +# so that, for instance, the PyTorch WDA solver works on an installation with no +# autograd or pymanopt. Functions raise an ImportError naming what they need. try: import autograd.numpy as np - from sklearn.decomposition import PCA + HAS_AUTOGRAD = True +except ImportError: # pragma: no cover - depends on the installation + import numpy as np + + HAS_AUTOGRAD = False + +try: import pymanopt import pymanopt.manifolds import pymanopt.optimizers -except ImportError: - raise ImportError( - "Missing dependency for ot.dr. Requires autograd, pymanopt, scikit-learn. You can install with install with 'pip install POT[dr]', or 'conda install autograd pymanopt scikit-learn'" - ) + + HAS_PYMANOPT = True +except ImportError: # pragma: no cover - depends on the installation + HAS_PYMANOPT = False + +try: + import torch + + HAS_TORCH = True +except ImportError: # pragma: no cover - depends on the installation + HAS_TORCH = False + +try: + from sklearn.decomposition import PCA + + HAS_SKLEARN = True +except ImportError: # pragma: no cover - depends on the installation + HAS_SKLEARN = False from .bregman import sinkhorn as sinkhorn_bregman from .utils import dist as dist_utils, check_random_state +def _require(condition, function, dependencies): + if not condition: + raise ImportError( + f"Missing dependency for ot.dr.{function}. Requires {dependencies}. " + "You can install with 'pip install POT[dr]', or " + "'conda install autograd pymanopt scikit-learn'" + ) + + def dist(x1, x2): r"""Compute squared euclidean distance between samples (autograd)""" x1p2 = np.sum(np.square(x1), 1) @@ -79,6 +111,184 @@ def split_classes(X, y): return [X[y == i, :].astype(np.float32) for i in lstsclass] +def _dist_torch(x1, x2): + r"""Squared euclidean distance between samples (torch).""" + return ( + torch.sum(x1**2, 1).reshape((-1, 1)) + + torch.sum(x2**2, 1).reshape((1, -1)) + - 2 * (x1 @ x2.T) + ) + + +def _sinkhorn_torch(w1, w2, M, reg, k): + r"""Sinkhorn algorithm with fixed number of iterations (torch).""" + K = torch.exp(-M / reg) + ui = torch.ones(M.shape[0], dtype=M.dtype, device=M.device) + vi = torch.ones(M.shape[1], dtype=M.dtype, device=M.device) + for _ in range(k): + vi = w2 / (K.T @ ui + 1e-50) + ui = w1 / (K @ vi + 1e-50) + return ui.reshape((-1, 1)) * K * vi.reshape((1, -1)) + + +def _sinkhorn_log_torch(w1, w2, M, reg, k): + r"""Sinkhorn algorithm in log-domain with fixed iterations (torch).""" + Mr = -M / reg + ui = torch.zeros(M.shape[0], dtype=M.dtype, device=M.device) + vi = torch.zeros(M.shape[1], dtype=M.dtype, device=M.device) + log_w1, log_w2 = torch.log(w1), torch.log(w2) + for _ in range(k): + vi = log_w2 - torch.logsumexp(Mr + ui[:, None], 0) + ui = log_w1 - torch.logsumexp(Mr + vi[None, :], 1) + return torch.exp(ui[:, None] + Mr + vi[None, :]) + + +def _stiefel_retract(P, X): + r"""QR retraction onto the Stiefel manifold, with a sign convention.""" + Q, R = torch.linalg.qr(P + X) + return Q * torch.sign(torch.sign(torch.diagonal(R)) + 0.5) + + +def _stiefel_project(P, G): + r"""Project a euclidean gradient onto the tangent space of Stiefel.""" + W = P.T @ G + return G - P @ (0.5 * (W + W.T)) + + +def _wda_cost_torch(P, xc, wc, regmean, reg, k, sinkhorn_solver): + r"""WDA objective: within-class transport cost over between-class.""" + loss_b, loss_w = 0.0, 0.0 + for i, xi in enumerate(xc): + xi = xi @ P + for j, xj in enumerate(xc[i:]): + xj = xj @ P + M = _dist_torch(xi, xj) + G = sinkhorn_solver(wc[i], wc[j + i], M, reg * regmean[i, j], k) + term = torch.sum(G * M) + if j == 0: + loss_w = loss_w + term + else: + loss_b = loss_b + term + if float(loss_b.detach() if torch.is_tensor(loss_b) else loss_b) == 0.0: + raise ValueError( + "The between-class transport cost underflowed to zero, so the WDA " + "objective is undefined. reg is too small for the scale of the " + "data: exp(-M/reg) underflows. Increase reg, or use " + "sinkhorn_method='sinkhorn_log'." + ) + return loss_w / loss_b + + +def _wda_torch(X, y, p, reg, k, sinkhorn_method, maxiter, verbose, P0, normalize): + r"""WDA solved with PyTorch autodiff and Riemannian gradient descent. + + Mirrors the pymanopt ``SteepestDescent`` path: projected gradient, QR + retraction and backtracking, so both solvers target the same optimum. + """ + dtype = X.dtype + device = X.device + labels = torch.unique(y) + xc = [X[y == c] for c in labels] + wc = [ + torch.full((x.shape[0],), 1.0 / x.shape[0], dtype=dtype, device=device) + for x in xc + ] + d = X.shape[1] + nc = len(xc) + + if P0 is None: + P = torch.linalg.qr(torch.randn(d, p, dtype=dtype, device=device))[0] + else: + P = P0.clone().to(dtype) + + regmean = torch.ones((nc, nc), dtype=dtype, device=device) + if P0 is not None and normalize: + with torch.no_grad(): + for i, xi in enumerate(xc): + xi = xi @ P + for j, xj in enumerate(xc[i:]): + xj = xj @ P + regmean[i, j] = torch.sum(_dist_torch(xi, xj)) / ( + xi.shape[0] * xj.shape[0] + ) + + if sinkhorn_method.lower() == "sinkhorn": + solver_fn = _sinkhorn_torch + elif sinkhorn_method.lower() == "sinkhorn_log": + solver_fn = _sinkhorn_log_torch + else: + raise ValueError("Unknown Sinkhorn method '%s'." % sinkhorn_method) + + def value(Q): + with torch.no_grad(): + return _wda_cost_torch(Q, xc, wc, regmean, reg, k, solver_fn) + + f = value(P) + step = 1.0 + for it in range(maxiter): + Q = P.detach().requires_grad_(True) + v = _wda_cost_torch(Q, xc, wc, regmean, reg, k, solver_fn) + (g,) = torch.autograd.grad(v, Q) + direction = -_stiefel_project(P, g) + gnorm = float(torch.linalg.norm(direction)) + if verbose: + print(f"{it + 1:<6d} {float(v.detach()):+.16e} {gnorm:.8e}") + if gnorm <= 1e-12: + break + # start from twice the last accepted step, as pymanopt's backtracking + # line search does, so progress is not throttled by a fixed unit step + step = min(2.0 * step, 1e4 / (gnorm + 1e-12)) + improved = False + for _ in range(40): + Pn = _stiefel_retract(P, step * direction) + fn = value(Pn) + if fn < f: + improved = True + break + step *= 0.5 + if not improved: + break + P, f = Pn, fn + return P.detach() + + +def _wda_torch_entry(X, y, p, reg, k, sinkhorn_method, maxiter, verbose, P0, normalize): + r"""Convert inputs, centre, run the torch solver, return ``(P, proj)``. + + numpy in gives numpy out; a torch tensor in keeps its device and dtype. + """ + was_numpy = not torch.is_tensor(X) + Xt = torch.as_tensor(X) if was_numpy else X + if not torch.is_floating_point(Xt): + Xt = Xt.to(torch.float64) + yt = y if torch.is_tensor(y) else torch.as_tensor(np.asarray(y)) + if P0 is None: + P0t = None + else: + P0t = (P0 if torch.is_tensor(P0) else torch.as_tensor(P0)).to(Xt.dtype) + + mx = Xt.mean(dim=0) + Xc = Xt - mx.reshape((1, -1)) + + Popt = _wda_torch( + Xc, yt, p, reg, k, sinkhorn_method, maxiter, verbose, P0t, normalize + ) + + if was_numpy: + Pn = Popt.detach().cpu().numpy() + mxn = mx.detach().cpu().numpy() + + def proj(Z): + return (Z - mxn.reshape((1, -1))).dot(Pn) + + return Pn, proj + + def proj(Z): + return (Z - mx.reshape((1, -1))) @ Popt + + return Popt, proj + + def fda(X, y, p=2, reg=1e-16): r"""Fisher Discriminant Analysis @@ -184,9 +394,17 @@ def wda( Size of dimensionality reduction. reg : float, optional Regularization term >0 (entropic regularization) - solver : None | str, optional - None for steepest descent or 'TrustRegions' for trust regions algorithm - else should be a pymanopt.solvers + solver : None | str | pymanopt.optimizers, optional + Chooses both the autodiff framework and the optimizer. + + - `None` or `'autograd'` (default): autograd and pymanopt + `SteepestDescent`. + - `'TrustRegions'` (or `'tr'`): autograd and pymanopt `TrustRegions`. + - a `pymanopt.optimizers` instance: autograd with that optimizer. + - `'torch'`: PyTorch autodiff with Riemannian gradient descent and a QR + retraction. Requires only `torch`, so it works on installations + without autograd or pymanopt, and accepts torch tensors directly, + keeping their device and dtype. sinkhorn_method : str method used for the Sinkhorn solver, either 'sinkhorn' or 'sinkhorn_log' P0 : ndarray, shape (d, p) @@ -211,6 +429,20 @@ def wda( Wasserstein Discriminant Analysis. arXiv preprint arXiv:1608.08063. """ # noqa + if solver == "torch": + _require(HAS_TORCH, "wda(solver='torch')", "torch") + return _wda_torch_entry( + X, y, p, reg, k, sinkhorn_method, maxiter, verbose, P0, normalize + ) + + _require( + HAS_AUTOGRAD and HAS_PYMANOPT, + "wda(solver='autograd')", + "autograd and pymanopt", + ) + if solver == "autograd": + solver = None + if sinkhorn_method.lower() == "sinkhorn": sinkhorn_solver = sinkhorn elif sinkhorn_method.lower() == "sinkhorn_log": @@ -495,6 +727,7 @@ def ewca( X = X - X.mean(0) if U0 is None: + _require(HAS_SKLEARN, "ewca", "scikit-learn") pca_fitted = PCA(n_components=k).fit(X) U = pca_fitted.components_.T if method == "MM": diff --git a/test/test_dr.py b/test/test_dr.py index dcb477717..7a7fdb92e 100644 --- a/test/test_dr.py +++ b/test/test_dr.py @@ -13,10 +13,17 @@ try: # test if autograd and pymanopt are installed import ot.dr - nogo = False + nogo = not (ot.dr.HAS_AUTOGRAD and ot.dr.HAS_PYMANOPT) except ImportError: nogo = True +try: + import torch + + notorch = False +except ImportError: + notorch = True + @pytest.mark.skipif(nogo, reason="Missing modules (autograd or pymanopt)") def test_fda(): @@ -251,3 +258,141 @@ def test_ewca(): U_last_eigvec = np.linalg.svd(X.T, full_matrices=False)[0][:, -k:] _, cos, _ = np.linalg.svd(U.T @ U_last_eigvec, full_matrices=False) assert np.allclose(cos, np.ones(k), atol=1e-3) + + +@pytest.mark.skipif(notorch, reason="Missing module (torch)") +def test_wda_torch_solver(): + rng = np.random.RandomState(0) + xs, ys = ot.datasets.make_data_classif("gaussrot", 90, random_state=rng) + xs = np.hstack((xs, rng.randn(90, 4))) + p = 2 + + P, proj = ot.dr.wda(xs, ys, p, maxiter=10, solver="torch") + + np.testing.assert_allclose(np.sum(P**2, 0), np.ones(p), rtol=1e-6) + assert proj(xs).shape == (90, p) + + +@pytest.mark.skipif(notorch, reason="Missing module (torch)") +def test_wda_torch_accepts_torch_tensors(): + rng = np.random.RandomState(0) + xs, ys = ot.datasets.make_data_classif("gaussrot", 90, random_state=rng) + xt = torch.tensor(xs, dtype=torch.float64) + yt = torch.tensor(ys) + + P, proj = ot.dr.wda(xt, yt, 2, maxiter=5, solver="torch") + + assert torch.is_tensor(P) + assert P.dtype == torch.float64 + assert proj(xt).shape == (90, 2) + + +@pytest.mark.skipif(notorch, reason="Missing module (torch)") +def test_wda_torch_does_not_modify_input(): + rng = np.random.RandomState(0) + xs, ys = ot.datasets.make_data_classif("gaussrot", 90, random_state=rng) + xs = xs + 10.0 + xs_copy = xs.copy() + + ot.dr.wda(xs, ys, 2, maxiter=5, solver="torch") + + np.testing.assert_allclose(xs, xs_copy) + + +@pytest.mark.skipif(notorch, reason="Missing module (torch)") +def test_wda_torch_sinkhorn_log(): + rng = np.random.RandomState(0) + xs, ys = ot.datasets.make_data_classif("gaussrot", 90, random_state=rng) + p = 2 + + P, _ = ot.dr.wda( + xs, ys, p, maxiter=10, solver="torch", sinkhorn_method="sinkhorn_log" + ) + + np.testing.assert_allclose(np.sum(P**2, 0), np.ones(p), rtol=1e-6) + + +@pytest.mark.skipif(nogo or notorch, reason="Missing modules") +def test_wda_backends_agree_on_cost_and_gradient(): + """The torch objective and its gradient must match the autograd ones.""" + import autograd + import autograd.numpy as anp + + rng = np.random.RandomState(0) + n, d, C, reg, k = 180, 6, 3, 1.0, 10 + X = np.vstack([rng.randn(n // C, d) + 3 * rng.randn(1, d) for _ in range(C)]) + y = np.repeat(np.arange(C), n // C) + P0 = np.linalg.qr(rng.randn(d, 2))[0] + Xc = X - X.mean(0) + + xc = [np.ascontiguousarray(Xc[y == c]) for c in range(C)] + wc = [np.ones(x.shape[0]) / x.shape[0] for x in xc] + + def cost_autograd(P): + loss_b, loss_w = 0.0, 0.0 + for i, xi in enumerate(xc): + xi = anp.dot(xi, P) + for j, xj in enumerate(xc[i:]): + xj = anp.dot(xj, P) + M = ot.dr.dist(xi, xj) + G = ot.dr.sinkhorn(wc[i], wc[j + i], M, reg, k) + term = anp.sum(G * M) + if j == 0: + loss_w = loss_w + term + else: + loss_b = loss_b + term + return loss_w / loss_b + + xct = [torch.tensor(x, dtype=torch.float64) for x in xc] + wct = [torch.tensor(w, dtype=torch.float64) for w in wc] + rmt = torch.ones((C, C), dtype=torch.float64) + Pt = torch.tensor(P0, dtype=torch.float64, requires_grad=True) + v = ot.dr._wda_cost_torch(Pt, xct, wct, rmt, reg, k, ot.dr._sinkhorn_torch) + (g,) = torch.autograd.grad(v, Pt) + + np.testing.assert_allclose(float(v.detach()), float(cost_autograd(P0)), rtol=1e-10) + np.testing.assert_allclose( + g.numpy(), autograd.grad(cost_autograd)(P0), rtol=1e-8, atol=1e-10 + ) + + +@pytest.mark.skipif(nogo or notorch, reason="Missing modules") +def test_wda_solvers_reach_comparable_objective(): + """Both solvers minimise the same objective, so neither should be much worse.""" + rng = np.random.RandomState(0) + xs, ys = ot.datasets.make_data_classif("gaussrot", 120, random_state=rng) + xs = np.hstack((xs, rng.randn(120, 3))) + P0 = np.linalg.qr(rng.randn(xs.shape[1], 2))[0] + + Pa, _ = ot.dr.wda(xs, ys, 2, maxiter=40, P0=P0) + Pt, _ = ot.dr.wda(xs, ys, 2, maxiter=40, P0=P0, solver="torch") + + Xc = torch.tensor(xs - xs.mean(0), dtype=torch.float64) + yt = torch.tensor(ys) + xc = [Xc[yt == c] for c in torch.unique(yt)] + wc = [torch.full((x.shape[0],), 1.0 / x.shape[0], dtype=torch.float64) for x in xc] + rm = torch.ones((len(xc), len(xc)), dtype=torch.float64) + + def objective(P): + return float( + ot.dr._wda_cost_torch( + torch.tensor(P, dtype=torch.float64), + xc, + wc, + rm, + 1, + 10, + ot.dr._sinkhorn_torch, + ) + ) + + assert objective(Pt) < 1.15 * objective(Pa) + + +@pytest.mark.skipif(notorch, reason="Missing module (torch)") +def test_wda_torch_unknown_sinkhorn_method(): + rng = np.random.RandomState(0) + xs, ys = ot.datasets.make_data_classif("gaussrot", 60, random_state=rng) + + with pytest.raises(ValueError): + ot.dr.wda(xs, ys, 2, maxiter=2, solver="torch", sinkhorn_method="nope") From 227d5f68614851936f1c88ffed71550c93b05ab3 Mon Sep 17 00:00:00 2001 From: Ahmed Eldeeb <62363199+deeb01@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:31:40 -0700 Subject: [PATCH 2/3] Correct PR number in RELEASES.md --- RELEASES.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/RELEASES.md b/RELEASES.md index 657bf11de..a3da914a0 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -4,8 +4,8 @@ #### New features -- `ot.dr.wda` gains `solver='torch'`, a PyTorch autodiff solver with Riemannian gradient descent, usable on installations without autograd or pymanopt, and accepting torch tensors directly (PR #853, Issue #806) -- `ot.dr` dependencies (autograd, pymanopt, scikit-learn, torch) are now imported optionally, so importing `ot.dr` no longer requires all of them; each function raises an explicit `ImportError` naming what it needs (PR #853) +- `ot.dr.wda` gains `solver='torch'`, a PyTorch autodiff solver with Riemannian gradient descent, usable on installations without autograd or pymanopt, and accepting torch tensors directly (PR #858, Issue #806) +- `ot.dr` dependencies (autograd, pymanopt, scikit-learn, torch) are now imported optionally, so importing `ot.dr` no longer requires all of them; each function raises an explicit `ImportError` naming what it needs (PR #858) ## 0.9.8dev From 522802ad6cf2ba2e27935177fbcc0941a5a4de2a Mon Sep 17 00:00:00 2001 From: Ahmed Eldeeb <62363199+deeb01@users.noreply.github.com> Date: Thu, 17 Sep 2026 11:32:06 -0700 Subject: [PATCH 3/3] Fix three differences between the torch and autograd wda solvers Found while reviewing the diff, all three cases where solver='torch' behaved differently from the default solver: - Non-numeric labels raised TypeError. numpy's split_classes indexes classes by value, so string labels work there, but torch.unique cannot hold them. Labels are now mapped to positional codes first. - p > d silently returned a wrongly shaped P. torch.linalg.qr on a (d, p) matrix with p > d returns a (d, d) factor, and nothing downstream objected. pymanopt's Stiefel(d, p) raises for this, so the torch path now checks 1 <= p <= d explicitly. - float32 numpy input returned float32 while the autograd path returns float64. numpy input is now promoted to float64; a torch tensor still keeps its own dtype, as documented. Adds a regression test for each. --- ot/dr.py | 18 ++++++++++++++---- test/test_dr.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/ot/dr.py b/ot/dr.py index a3922ee22..e55f92115 100644 --- a/ot/dr.py +++ b/ot/dr.py @@ -196,6 +196,9 @@ def _wda_torch(X, y, p, reg, k, sinkhorn_method, maxiter, verbose, P0, normalize d = X.shape[1] nc = len(xc) + if not 1 <= p <= d: + raise ValueError(f"Need d >= p >= 1. Values supplied were d = {d} and p = {p}") + if P0 is None: P = torch.linalg.qr(torch.randn(d, p, dtype=dtype, device=device))[0] else: @@ -258,10 +261,17 @@ def _wda_torch_entry(X, y, p, reg, k, sinkhorn_method, maxiter, verbose, P0, nor numpy in gives numpy out; a torch tensor in keeps its device and dtype. """ was_numpy = not torch.is_tensor(X) - Xt = torch.as_tensor(X) if was_numpy else X - if not torch.is_floating_point(Xt): - Xt = Xt.to(torch.float64) - yt = y if torch.is_tensor(y) else torch.as_tensor(np.asarray(y)) + if was_numpy: + # match the autograd path, which promotes to float64 via P + Xt = torch.as_tensor(np.asarray(X, dtype=np.float64)) + else: + Xt = X if torch.is_floating_point(X) else X.to(torch.float64) + # labels may be strings or any hashable, which torch cannot hold, so index + # them by position the way numpy's split_classes does + if torch.is_tensor(y): + yt = y + else: + yt = torch.as_tensor(np.unique(np.asarray(y), return_inverse=True)[1]) if P0 is None: P0t = None else: diff --git a/test/test_dr.py b/test/test_dr.py index 7a7fdb92e..86ae9ce32 100644 --- a/test/test_dr.py +++ b/test/test_dr.py @@ -396,3 +396,35 @@ def test_wda_torch_unknown_sinkhorn_method(): with pytest.raises(ValueError): ot.dr.wda(xs, ys, 2, maxiter=2, solver="torch", sinkhorn_method="nope") + + +@pytest.mark.skipif(notorch, reason="Missing module (torch)") +def test_wda_torch_non_numeric_labels(): + """Labels need not be numeric: numpy indexes by value, torch cannot.""" + rng = np.random.RandomState(0) + xs = np.vstack([rng.randn(40, 5) + 3 * rng.randn(1, 5) for _ in range(2)]) + ys = np.array(["cat"] * 40 + ["dog"] * 40) + + P, _ = ot.dr.wda(xs, ys, 2, maxiter=3, solver="torch") + + assert P.shape == (5, 2) + + +@pytest.mark.skipif(notorch, reason="Missing module (torch)") +def test_wda_torch_rejects_p_larger_than_d(): + rng = np.random.RandomState(0) + xs, ys = ot.datasets.make_data_classif("gaussrot", 60, random_state=rng) + + with pytest.raises(ValueError): + ot.dr.wda(xs, ys, xs.shape[1] + 1, maxiter=3, solver="torch") + + +@pytest.mark.skipif(notorch, reason="Missing module (torch)") +def test_wda_torch_numpy_input_returns_float64(): + """numpy in gives float64 out, matching the autograd solver.""" + rng = np.random.RandomState(0) + xs, ys = ot.datasets.make_data_classif("gaussrot", 60, random_state=rng) + + P, _ = ot.dr.wda(xs.astype(np.float32), ys, 2, maxiter=3, solver="torch") + + assert P.dtype == np.float64