Add associated Legendre functions to JAX

[tlu7]

We have an implementation of associated Legendre functions of the first kind for spherical harmonics and spherical harmonics transforms. It is scipy equivalence is scipy.special.lpmn(m, n, z), where m is the order, n is the degree, and m <= n. The output array has shape (m+1, n+1, num_z). In our implementation, we have only n and z as the arguments by enforcing m=n such that the output array has shape (n+1, n+1, num_z). By enforcing m=n, our implementation utilizes three recurrence, which is faster. In addition, for applications related to spherical harmonics and spherical harmonic transform, people usually calculate all orders m up to the maximum degree n.

We would like to have feedbacks from the community on how to contribute this work in JAX:
option 1: following the input pattern of scipy.speical.lpmn(m, n, z) and raise error if m!=n
option 2: making it a standalone function lpnn(n, z)
and option 3: your ideas

[tlu7]

Looping in @hawkinsp @shoyer

[shoyer]

JAX tries not to deviate from NumPy or SciPy APIs unless absolutely necessary -- and if it is necessary, we typically try to come up new primitives (e.g., in jax.lax) rather than deviating from scipy in the jax.scipy namespace. So in this case, I would pick your first option.

[tlu7]

Thanks for the suggestion!

[jakevdp]

In your implementation m and n are static, correct? In that case, I think defining jax.scipy.speical.lpmn with a NotImplementedError for m != n is a good way to go.

[tlu7]

Thanks for the suggestion! Yes, both m and n are static.