Add JAX Support to NPBench and Implement JAX Benchmarks

[hardik01shah]

Overview:

This PR introduces JAX as a supported framework in NPBench, providing JAX implementations for all existing benchmarks. This enhancement allows NPBench to compare the performance of JAX against other frameworks, for all implemented benchmarks in NPBench.

---

Changes Introduced:

1. Framework Configuration:
   • Added framework_info/jax.json to define framework-specific configuration settings for JAX.
2. Benchmark Implementations:
   • Added JAX implementations for all benchmarks in the directory structure:
     • npbench/benchmarks/<bench_parent>/<bench_name>_jax.py.
   • Support for benchmarks with existing JAX library implementations.
     • npbench/benchmarks/<bench_parent>/<bench_name>_jax_lib.py
3. Infrastructure Updates:
   • Updated npbench/infrastructure/__init__.py to include JAX framework initialization.
   • Added npbench/infrastructure/jax_framework.py to handle:
     • Pre- and post-processing for JAX array arguments (e.g., using block_until_ready() and jnp.array()).
4. Bug Fix:
   • Resolved a minor bug in jax_framework.py related to incorrect naming of the current framework during benchmarking.

---

Motivation:

The addition of JAX support enhances NPBench by:

• Allowing direct comparison of JAX’s performance with other frameworks across diverse scientific Python benchmarks.
• Extending the utility of NPBench to the JAX community for performance evaluation and optimization.

---

Testing and Validation:

• Verified the correctness of JAX implementations by validating the implementations with the Numpy impplementations.
• Ensured smooth integration of JAX into the benchmarking pipeline.

---

Contributors:

• Hardik Shah
• Sushant Swamy
• Filip Jaksic