From 35c0f64836adee9ab71a6d06cc6303226da684bf Mon Sep 17 00:00:00 2001 From: Jake VanderPlas Date: Tue, 6 Feb 2024 11:37:42 -0800 Subject: [PATCH] jnp.linalg.solve: deprecate batched 1D solves when b.ndim > 1 --- CHANGELOG.md | 3 +++ jax/_src/numpy/linalg.py | 13 ++++++++++--- tests/linalg_test.py | 2 ++ 3 files changed, 15 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a4af11e7f72f..486998b6321b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -110,6 +110,9 @@ Remember to align the itemized text with the first line of an item within a list such as {func}`~jax.random.bits`, {func}`~jax.random.gamma`, and others, is deprecated and will emit a `FutureWarning`. Use `jax.vmap` for explicit batching. * {func}`jax.lax.tie_in` is deprecated: it has been a no-op since JAX v0.2.0. + * {func}`jax.numpy.linalg.solve` now shows a deprecation warning for batched 1D + solves with `b.ndim > 1`. In the future these will be treated as batched 2D + solves. ## jaxlib 0.4.24 diff --git a/jax/_src/numpy/linalg.py b/jax/_src/numpy/linalg.py index cf1b1c9b5eb3..65ba2c9105a2 100644 --- a/jax/_src/numpy/linalg.py +++ b/jax/_src/numpy/linalg.py @@ -16,6 +16,7 @@ from collections.abc import Sequence from functools import partial +import warnings import numpy as np import textwrap @@ -635,9 +636,15 @@ def qr(a: ArrayLike, mode: str = "reduced") -> Array | QRResult: def solve(a: ArrayLike, b: ArrayLike) -> Array: check_arraylike("jnp.linalg.solve", a, b) a, b = promote_dtypes_inexact(jnp.asarray(a), jnp.asarray(b)) - # TODO(jakevdp): this condition matches the broadcasting behavior in numpy < 2.0. - # For the array API specification, we would check only if b.ndim == 1. - if b.ndim == 1 or a.ndim == b.ndim + 1: + + if b.ndim == 1: + signature = "(m,m),(m)->(m)" + elif a.ndim == b.ndim + 1: + # Deprecation warning added 2024-02-06 + warnings.warn("jnp.linalg.solve: batched 1D solves with b.ndim > 1 are deprecated, " + "and in the future will be treated as a batched 2D solve. " + "Use solve(a, b[..., None])[..., 0] to avoid this warning.", + category=FutureWarning) signature = "(m,m),(m)->(m)" else: signature = "(m,m),(m,n)->(m,n)" diff --git a/tests/linalg_test.py b/tests/linalg_test.py index d62fea067cc7..7e0e93942e0d 100644 --- a/tests/linalg_test.py +++ b/tests/linalg_test.py @@ -1035,6 +1035,7 @@ def testSolve(self, lhs_shape, rhs_shape, dtype): lhs_shape=[(2, 2), (2, 2, 2), (2, 2, 2, 2), (2, 2, 2, 2, 2)], rhs_shape=[(2,), (2, 2), (2, 2, 2), (2, 2, 2, 2)] ) + @jtu.ignore_warning(category=FutureWarning, message="jnp.linalg.solve: batched") def testSolveBroadcasting(self, lhs_shape, rhs_shape): # Batched solve can involve some ambiguities; this test checks # that we match NumPy's convention in all cases. @@ -1196,6 +1197,7 @@ def test(x): self.assertAllClose(xc, grad_test_jc(xc)) @jtu.skip_on_flag("jax_skip_slow_tests", True) + @jtu.ignore_warning(category=FutureWarning, message="jnp.linalg.solve: batched") def testIssue1151(self): rng = self.rng() A = jnp.array(rng.randn(100, 3, 3), dtype=jnp.float32)