Providing pre-built TensorFlow installation?

[deniskristak]

Hello,
I'm trying to add support of jax to our HPC package installator/dependency manager, Easybuild.
I would like to know if it would be appropriate to provide TensorFlow, which is already built, to jax as dependency, without bazel building it again on it's own.
We do have support for TensorFlow already so that would seem like the most logical thing to do here (managing deps in one installation SW is more practical than to think about Bazel's management as well).
If so, how should I do it (not much experience with Bazel)?
Should I follow this tutorial - non-bazel-projects part?
Another POI is whether it is safe (doesn't corrupt the TensorFlow installation). I've seen here that jax might make some changes to TensorFlow installation?
Thank you very much for helping us out, have a great day! :-)

[hawkinsp]

jax does not depend on TensorFlow itself. jax depends on XLA, which happens to be bundled as part of the TensorFlow tree. The tensorflow package that is distributed does not expose the XLA APIs we would need even if we wanted to do this. So no, I don't think this is a good idea.

(We might more realistically hope for XLA to be distributed separately from TF, but there's no concrete plan to do that.)

[deniskristak]

Thanks for a quick clarification, how can I Install jax, when we already have TF with XLA installed and optimized

[hawkinsp]

I'd just install JAX using the standard wheels (https://github.com/google/jax#installation).

If for some reason you want to build from source, you can do so with https://jax.readthedocs.io/en/latest/developer.html#building-from-source

But note that I would not expect a significant performance improvement from building from source. JAX is mostly a jit compiler, and generates code bespoke to the machine it is running on.

Having TF built or installed does not help with building or installing JAX. They have only a source-level dependency relationship; the deployed packages are almost entirely separate.

[boegel]

It's not only about performance when building from source. It's also about supporting non-x86_64 systems (like Arm and POWER). And making sure that the (parts of) TensorFlow that are being used come from the same version.

When building from source, we also strongly prefer to not having anything downloaded on the fly underneath the covers.

Is there any support currently for pointing to an existing TensorFlow installation when installing jax(lib)?

[boegel]

I meant not having anything downloaded on the fly, of course (edited my previous comment)

[hawkinsp]

As I mentioned above: there's no benefit to having a TensorFlow binary installed. JAX cannot use TensorFlow binary distributions and nor do we ever plan to do that. JAX incorporates the pieces of the TF source tree it needs (XLA) in a separate module (jaxlib). TF does not expose what we need and is not released on a schedule that is useful for JAX. The copy of XLA cannot be shared in binary form between JAX and TF and nor is there any benefit from doing so.

Further, there is no requirement that JAX and TF, if both installed, have matching versions. Each bundles the things they need internally: they do not conflict or interact. Even though they share certain components, each bundles their own private copy.

You should be able to pre-download dependencies from source when building with Bazel. If you have a way of building TF you are happy with that solves this, you will probably find that you can do the same things for building jaxlib (e.g., point Bazel to precached tarballs for dependencies instead of having Bazel fetch it.)