[D] JAX learning resources?

This page summarizes the projects mentioned and recommended in the original post on /r/JAX

Our great sponsors
  • WorkOS - The modern identity platform for B2B SaaS
  • InfluxDB - Power Real-Time Data Analytics at Scale
  • SaaSHub - Software Alternatives and Reviews
  • flax

    Flax is a neural network library for JAX that is designed for flexibility.

    - https://github.com/google/flax/tree/main/examples

  • dm-haiku

    JAX-based neural network library

    - https://github.com/deepmind/dm-haiku/tree/main/examples

  • WorkOS

    The modern identity platform for B2B SaaS. The APIs are flexible and easy-to-use, supporting authentication, user identity, and complex enterprise features like SSO and SCIM provisioning.

  • jaxline

    - https://github.com/deepmind/jaxline (training loop code used by below)

  • deepmind-research

    This repository contains implementations and illustrative code to accompany DeepMind publications

    - https://github.com/deepmind/deepmind-research (i read `nfnets` and `perceiver`, there are more there)

NOTE: The number of mentions on this list indicates mentions on common posts plus user suggested alternatives. Hence, a higher number means a more popular project.

Suggest a related project

Related posts