Make JAX optional

[xiki-tempula]

Pymbar is making JAX optional with pymbar-core (no JAX) and pymbar (with JAX)
choderalab/pymbar#500

Do we want to have a similar arrangement of alchemlyb-core and alchemlyb?

[orbeckst]

We could, although until someone really needs the slimmer, slower *-core version, I'd say we use high-performance by default where possible.

[orbeckst]

Note issue #336 regarding JAX issues.

[orbeckst]

With conda-forge/alchemlyb-feedstock#25 we will have conda-forge packages alchemlyb (includes pymbar and thus jax) and alchemlyb-core (depends on pymbar-core and no jax).

We will just need to document these choices in the installation instructions https://alchemlyb.readthedocs.io/en/latest/install.html .

[xiki-tempula]

@orbeckst Sorry we are not able to make the alchemlyb (includes pymbar and thus jax) and alchemlyb-core (depends on pymbar-core and no jax) to work.
My attempt: conda-forge/alchemlyb-feedstock#24
Attempt by another guy: conda-forge/alchemlyb-feedstock#26

The solution that I came up with is to use build variant
By default it will pull the jax build varient, but you could get the non-JAX version with build string selector alchemlyb=*=core_*
In this PR conda-forge/alchemlyb-feedstock#25

[orbeckst]

Thanks for the explanation and clarification @xiki-tempula – I had only briefly skim-read the conda-forge PRs. We should still document the installation options with variants as they will solve the problem raised in this issue.

[xiki-tempula]

@orbeckst Are you happy for me to merge the conda-forge PR? We could then change the documentation.

[orbeckst]

Of course, yes please.

…

On Aug 28, 2025, at 12:59, Zhiyi Wu ***@***.***> wrote: xiki-tempula left a comment (alchemistry/alchemlyb#327) <#327 (comment)> @orbeckst <https://github.com/orbeckst> Are you happy for me to merge the conda-forge PR? We could then change the documentation. — Reply to this email directly, view it on GitHub <#327 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAB2DHE7KK22APF5BQN35BT3P5NRTAVCNFSM6AAAAACFCI5FA2VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZTEMZUG43TAMBVGE>. You are receiving this because you were mentioned.

-- Oliver Beckstein (he/his/him) ***@***.***

[orbeckst]

@xiki-tempula , given that we just merged PR #437 for handling jax deps, can we close this issue as resolved as there is a solution?