JAX

JAX

Autograd and XLA

• GitHub
• Doc: Autodiff Cookbook
• Try JAX on Cloud TPUs in Colab
• Talk by Skye Wanderman-Milne
  • Slides: Accelerated machine-learning researchvia composable function transformationsin Python

What is JAX?

• Numpy on GPU
  • Leverages XLA compiler
    • XLA (Accelerated Linear Algebra) is a domain-specific compiler for linear algebra that can accelerate TensorFlow models with potentially no source code changes.
• Autodifferentiation capability
  • How to get gradients of some loss function your code computes with respect to you input parameters

What is Elegy?

A higher level of JAX.

• Elegy

Read more

• Plugging Into JAX

DL Frameworks

• From PyTorch to JAX: towards neural net frameworks that purify stateful code

Jax vs PyTorch

• ← Previous
Next →