Modernize JAX compatibility and separate core dependencies from example extras

[vimarsh6739]

Hi there! I'm working on adding jaxpi as a benchmark in Enzyme-JAX, a compiler that applies various compiler optimizations on top of the existing JAX pipeline.

I thought it'd be super cool to include JAX-PI in our CI benchmark suite, which runs on every Enzyme-JAX commit and publishes it's results here

I had to make the following changes:

• Update the JAX API in a few places (I tested the changes locally with JAX 0.11)
• Re-organize the packaging into core + example-specific extras
• Remove the pytorch dependency from the sampler base class (it seemed unused)

Specifically, this allows us to install and import jaxpi without pulling in the visualization and experiment-tracking dependencies.

Hopefully, these changes are useful to you all too!

[vimarsh6739]

For the initial Enzyme-JAX integration, I'm temporarily pinning the dependency to a commit in my fork. If these changes are accepted, I'd be happy to update Enzyme-JAX to track this upstream repository directly instead. cc @sifanexisted