Bug#1143497: RFP: python-jax -- Differentiate, compile, and transform Numpy code

Package: wnpp
Severity: wishlist
X-Debbugs-Cc: [email protected]

* Package name    : python-jax
  Version         : 0.11.0
* URL             : https://github.com/jax-ml/jax
* License         : Apache 2.0
  Programming Lang: C++, Python
  Description     : Differentiate, compile, and transform Numpy code

JAX is a Python library for accelerator-oriented array computation and
program transformation, designed for high-performance numerical computing
and large-scale machine learning.

JAX can automatically differentiate native Python and NumPy functions.
It can differentiate through loops, branches, recursion, and closures,
and it can take derivatives of derivatives of derivatives. It supports
reverse-mode differentiation (a.k.a. backpropagation) via `jax.grad` as
well as forward-mode differentiation, and the two can be composed arbitrarily
to any order.

JAX uses XLA to compile and scale your NumPy programs on TPUs, GPUs, and
other hardware accelerators. You can compile your own pure functions with
`jax.jit`. Compilation and automatic differentiation can be composed arbitrarily.


This package is an (optional) dependency for the refnx package but will
be useful to other packages as a way of speeding up numerical
calculations from Python.

The build system for jaxlib is bazel - this will take some wrangling.

OpenXLA also needs to be packaged for jaxlib. https://github.com/openxla/xla