Недавно Google Cloud выпустил «Руководство разработчика PyTorch по основам JAX».

Jax – это фреймворк для машинного обучения, подобный PyTorch и TensorFlow.

Его разработали в Deepmind, хотя он не является официальным продуктом Google, он остается популярным.

Jax объединяет Autograd и XLA (Accelerated Linear Algebra - компилятор с открытым исходным кодом для машинного обучения) для обеспечения высокопроизводительных численных вычислений.

Созданный на основе NumPy, его синтаксис следует той же структуре, что делает его простым выбором для разработчиков.

В этом руководстве содержится пошаговый гайд по  реализации простой нейтронной сети на Pytorch (JAX + Flax NNX) для тех, кто хочет начать работать с JAX.