Remove JAX implementation

[jwilles]

The project currently includes a JAX implementation that is no longer needed or maintained. Please remove all JAX-related code, dependencies, and documentation from the repository to streamline the codebase and reduce maintenance overhead.

Tasks:

• Identify all files and modules associated with JAX.
• Remove JAX code and any related scripts or utilities.
• Update documentation and requirements files to reflect the removal of JAX.
• Ensure tests and CI are updated accordingly.

This cleanup will help focus development efforts on supported backends and simplify maintenance.