smolyax: a high-performance implementation of the Smolyak interpolation operator in JAX

The smolyax library provides interpolation capabilities for arbitrary multivariate and vectorvalued functions in out for any in , out .It implements the Smolyak interpolation operator, which is known to overcome the curse-of-dimensionality plaguing naive multivariate interpolation (Barthelmann et al., 2000) and uses the barycentric interpolation formula (Berrut & Trefethen, 2004) for numerical stability.The implementation is based on JAX (Bradbury et al., 2018), a free and open-source Python library for high-performance computing that integrates with the Python and NumPy numerical computing ecosystem.Thanks to JAX's device management, smolyax runs natively on both CPU and GPU.While implementing Smolyak interpolation in JAX is challenging due to the highly irregular data structures involved, smolyax overcomes this by employing a tailored batching and padding strategy that enables efficient vectorization.smolyax supports interpolation on bounded or unbounded domains via Leja (M. A. Chkifa, 2013) or Gauss-Hermite (Abramowitz & Stegun, 1964) node sequences, respectively.It provides efficient Numba-accelerated routines to generate isotropic or anisotropic total degree multi-index sets (Adcock et al., 2022, sec.2.3.2),which are the key ingredient to generate the high-dimensional sparse grids (Bungartz & Griebel, 2004) of interpolation nodes required by the Smolyak interpolation operator.Additional types of node sequences or multi-index sets can be incorporated easily by implementing a minimalistic interface.The smolyax interpolant provides further functionality to evaluate its gradient as well as compute its integral.