To build and install jax, jaxlib, pjrt and plugin at once.

[i-chaochen]

Motivation

Now we can just use one line to build everything (jax, jaxlib, pjrt and plugin) as we're used to.

python3 stack.py build --xla-dir=/my/own/xla/path --jax-dir=/my/own/jax/path

[charleshofer]

Looks fine to me

[charleshofer]

Could you rebase and retarget this for master?

[i-chaochen]

@charleshofer this is master one #244 , I created two PRs

[Arech8]

Please don't merge until comments to #244 are addressed.

Also I'm not sure we should update released branch... Though this only affects the build system, so we could do it, just let's finish with the corresponding PR into the master first.

[Arech8]

Also why do you make two separate PRs (into the master and here), instead of having just one main PR merged into the master and then cherry-picking it here? I think, this PR should be discarded in favor of cherry-picking, this helps with maintenance.

[i-chaochen]

Please don't merge until comments to #244 are addressed.

Also I'm not sure we should update released branch... Though this only affects the build system, so we could do it, just let's finish with the corresponding PR into the master first.

Also why do you make two separate PRs (into the master and here), instead of having just one main PR merged into the master and then cherry-picking it here? I think, this PR should be discarded in favor of cherry-picking, this helps with maintenance.

Yes, I can close this one till #244 is addressed and cherry-pick back.