Hi, tech lead for TFP here. The wording here was unclear -- sorry! We're fixing it presently.
We are not migrating away from TF; far from it!
The change here was to interoperate with TF and JAX (and numpy!), by way of some rewrite trickery under the hood. Essentially, we wrote a translation layer that implements the TF API surface (or, the parts we actually use) in terms of numpy & JAX primitives [1]. This lets us leave most TFP code intact, written in terms of the TF API, but interoperate with JAX by way of the API translation layer. (Actually we implemented numpy support first, and mostly got JAX for "free" since JAX is largely API-compatible with numpy).
Sorry for any confusion!
We're pretty stoked about this work, so happy to answer any other questions you may have (also feel free to chime in on the github tracker or email [email protected])
We are not migrating away from TF; far from it!
The change here was to interoperate with TF and JAX (and numpy!), by way of some rewrite trickery under the hood. Essentially, we wrote a translation layer that implements the TF API surface (or, the parts we actually use) in terms of numpy & JAX primitives [1]. This lets us leave most TFP code intact, written in terms of the TF API, but interoperate with JAX by way of the API translation layer. (Actually we implemented numpy support first, and mostly got JAX for "free" since JAX is largely API-compatible with numpy).
Sorry for any confusion!
We're pretty stoked about this work, so happy to answer any other questions you may have (also feel free to chime in on the github tracker or email [email protected])
[1] - https://github.com/tensorflow/probability/tree/master/tensor...