Module: tff.jax

Libraries for interacting with a JAX frontend and XLA backend.

Functions

computation(...): Decorates/wraps Python functions containing JAX code as TFF computations.