Add Schechter distributions for astrophysical luminosity functions

Feature Summary

Introduce SchechterMag and DoubleSchechterMag distribution classes into NumPyro.

• The Schechter function is the standard parametric form used to describe galaxy luminosity and stellar mass functions.
• Both single and double forms are required in astronomy, with parameters (α, M*, φ*) (single) or two such sets (double).
• Key feature: implementation preserves JAX auto-differentiability, allowing the No-U-Turn Sampler (NUTS) to be applied in hierarchical Bayesian models.
• Our prototype implementations already exist in a standalone package (numpyro-schechter), but native NumPyro support would benefit a broader community.

See github and PyPI for the existing implementation.

Why is this needed?

• Astronomy and cosmology rely heavily on Bayesian hierarchical modeling to study galaxy populations. Current analyses are limited by the absence of a Schechter distribution in PyMC/NumPyro.
• Existing Python packages (e.g., scipy, sympy) can evaluate the Schechter function and integrals, but they are not JAX-compatible and cannot be differentiated automatically.
• A JAX-native Schechter implementation would:
  • Enable efficient gradient-based MCMC (NUTS) for models with hundreds of parameters.
  • Open NumPyro to a large class of astrophysical inference problems (luminosity/mass functions, environmental dependencies, multi-wavelength studies).
  • Provide a “plug-and-play” likelihood component for researchers beyond astronomy who may encounter similar truncated power-law distributions with exponential cutoffs.

Our group has already implemented these distributions externally and confirmed they work with NumPyro’s inference algorithms. We would like to contribute these back to NumPyro, with tests and documentation, so that they are maintained and available to the wider community.

Start by reviewing the existing implementation in the linked numpyro-schechter package and compare its SchechterMag and DoubleSchechterMag behavior with NumPyro's distribution conventions. Add both parameterizations with JAX autodifferentiability, then provide the requested tests and documentation demonstrating compatibility with NumPyro inference algorithms.