Optax is a gradient processing and optimization library for JAX. It is designed to facilitate research by providing building blocks that can be recombined in custom ways in order to optimize parametric models such as, but not limited to, deep neural networks. We favor focusing on small composable building blocks that can be effectively combined into custom solutions. Others may build upon these basic components in more complicated abstractions. Whenever reasonable, implementations prioritize readability and structuring code to match standard equations, over code reuse.

Features

  • Optax is a gradient processing and optimization library for JAX
  • Provides simple, well-tested, efficient implementations of core components
  • Documentation available
  • Examples available
  • Improve research productivity by enabling to easily combine low-level ingredients into custom optimisers (or other gradient processing components)
  • Accelerate adoption of new ideas by making it easy for anyone to contribute
  • You can install the latest released version of Optax from PyPI

Project Samples

Project Activity

See All Activity >

Categories

Machine Learning

License

Apache License V2.0

Follow Optax

Optax Web Site

You Might Also Like
Find out just how much your login box can do for your customer | Auth0 Icon
Find out just how much your login box can do for your customer | Auth0

With over 53 social login options, you can fast-track the signup and login experience for users.

From improving customer experience through seamless sign-on to making MFA as easy as a click of a button – your login box must find the right balance between user convenience, privacy and security.
Rate This Project
Login To Rate This Project

User Reviews

Be the first to post a review of Optax!

Additional Project Details

Operating Systems

Linux, Mac, Windows

Programming Language

Python

Related Categories

Python Machine Learning Software

Registered

5 days ago