GH-44229: [Docs] Add PyArrow to JAX example to the docs

[felipecrv]

Rationale for this change

Explicitly mention in the docs a way that PyArrow can interop with JAX.

What changes are included in this PR?

• Tweaks to the phrasing
• Two JAX examples: one for jax.numpy and another for jax.dlpack

Are these changes tested?

N/A

• GitHub Issue: [Python][Docs] Mention Arrow->JAX conversion via dlpack in the examples #44229

[github-actions]

⚠️ GitHub issue #44229 has been automatically assigned in GitHub to PR creator.

[amoeba]

LGTM.

One question, jax.numpy.from_dlpack ("Returns a Array representation of a DLPack tensor") and jax.dlpack.from_dlpack ("Construct a JAX array via DLPack") seem like redundant APIs, do JAX users tend to understand which is which or could we make it clear why a user would pick one or other in this example?

[felipecrv]

One question, jax.numpy.from_dlpack ("Returns a Array representation of a DLPack tensor") and jax.dlpack.from_dlpack ("Construct a JAX array via DLPack") seem like redundant APIs, do JAX users tend to understand which is which or could we make it clear why a user would pick one or other in this example?

Yes. JAX itself is a nice to use Python library wrapping the XLA compiler. Then it provides public interfaces inspired by famous Python libraries: there is jax.numpy and jax.scipy for instance. jax.dlpack is the standalone entrypoint to dlpack code in case you are not using JAX via the familiar interfaces.

[amoeba]

Got it, thanks.

[conbench-apache-arrow]

After merging your PR, Conbench analyzed the 4 benchmarking runs that have been run so far on merge-commit 9fa78d0.

There were no benchmark performance regressions. 🎉

The full Conbench report has more details. It also includes information about 1 possible false positive for unstable benchmarks that are known to sometimes produce them.

[AlenkaF]

Great addition, thanks @felipecrv!