sph_harm_y returns incorrect values for batched angles (same root cause as #20769)

Description

jax.scipy.special.sph_harm_y returns incorrect values whenever it is called with a batch of angles and a length-1 degree/order — the result is correct at index 0 and silently wrong everywhere else.

This is the same root cause as #20769 (which reports it through the now-deprecated sph_harm), but it is worth a separate report because it is still present in sph_harm_y, the API scipy 1.15+ and the JAX docs point users to, and because two further problems come along with it. Verified on jax 0.11.1 against scipy 1.18.1.

The mispairing

_sph_harm indexes its internal Legendre table positionally (jax/_src/scipy/special.py):

legendre_val = legendre.at[abs(m), n, jnp.arange(len(n))].get(mode="clip")

jnp.arange(len(n)) walks the angle axis with indices taken from the length of n, pairing n[i] with theta[i] positionally instead of broadcasting the four arguments against one another. A length-1 degree against a batch of angles therefore gathers the Legendre value at theta[0] only, and then multiplies that single value by the azimuthal factor at every phi.

Reproducer

import os
os.environ["JAX_ENABLE_X64"] = "True"

import jax.numpy as jnp
import jax.scipy.special as jss
import numpy as np
import scipy.special as sps

theta = jnp.linspace(0.1, 3.0, 4)
phi = jnp.zeros(4)

got = jss.sph_harm_y(jnp.array([3]), jnp.array([0]), theta, phi)
want = sps.sph_harm_y(3, 0, np.asarray(theta), np.asarray(phi))
print(np.asarray(got))
print(want)
print(np.abs(np.asarray(got) - want))

[0.72412021+0.j 0.72412021+0.j 0.72412021+0.j 0.72412021+0.j]
[ 0.72412021+0.j -0.3304791 +0.j  0.33377659+0.j -0.70209664+0.j]
[2.22044605e-16 1.05459931e+00 3.90343624e-01 1.42621685e+00]

Every entry is the value at theta[0]. The error reaches 1.43 absolute for n = 3, m = 0 — an ordinary call, not an edge case. Sweeping n <= 3 and |m| <= n, that is the worst case for these angles; the exact maximum depends on which angles are sampled, hence the explicit reproducer.

Rank-0 input

len(n) is called on the degree directly, so a scalar degree does not fail gracefully:

jss.sph_harm_y(3, 0, theta, phi)
# TypeError: len() of unsized object

Derivatives at the poles

Separately, the derivative is nan at both poles for every n >= 1, because the implementation differentiates sqrt(1 - cos^2 theta) and the surviving 0 * inf is nan:

n=0: grad(theta=0)=0.0  grad(theta=pi)=0.0
n=1: grad(theta=0)=nan  grad(theta=pi)=nan
n=2: grad(theta=0)=nan  grad(theta=pi)=nan
n=3: grad(theta=0)=nan  grad(theta=pi)=nan

Notes

This reached released downstream code without any test catching it, because every fixture evaluated a single position — the one index that is right. A suite that only ever checks one angle cannot see this class of bug, which may be worth keeping in mind for whatever fix lands.

For what it's worth, on the fix direction: taking the degree and order as static Python ints leaves nothing to mispair, and building sin^m(theta) as an integer power never forms the square root, so the pole derivatives stay finite.

I'm aware of the position in #20769 that this code may be removed rather than fixed (scope JEP). If that is still the plan, a deprecation warning on sph_harm_y would be worth more than silence — right now it is the recommended spelling and it returns wrong numbers without complaint.

System info (python version, jaxlib version, accelerator, etc.)

jax:        0.11.1
jaxlib:     0.11.1
scipy:      1.18.1
python:     3.14.3
platform:   macOS-26.6.2-arm64
device:     CpuDevice(id=0)