PaperlayerView exact revision on arXiv

Source-linked AI summary

Equinox: neural networks in JAX via callable PyTrees and filtered transformations

Patrick Kidger, Cristian Garcia

arXiv:2111.00254v1cs.LGcs.PL

TL;DR

JAX’s pure-function model and PyTorch’s class-based model-building syntax create a design tension for parameterised functions. Equinox represents those functions as callable PyTrees and filters PyTrees at transformation boundaries, demonstrating class-based modeling without sacrificing JAX-style functional programming. The authors report that this avoids the abstractions, conversions, and integration restrictions associated with earlier libraries.

  • Problem

    JAX libraries have either rejected OO syntax or introduced transformations, abstractions, and integration limitations to support class-based parameterised functions.

  • Method

    Equinox represents parameterised functions as callable PyTrees and filters PyTrees to isolate components for jit, grad, or vmap transformations.

  • Results

    Equinox demonstrates that PyTorch-like class-based model syntax can be used without sacrificing JAX-style functional programming.

  • Takeaways & Limitations

    Using only PyTrees and transformations, Equinox avoids new abstractions, OO-to-functional translation, and special interoperability requirements.

  • Takeaways & Limitations

    The simple parameterised-function example assumes every PyTree leaf has JAX differentiation rules; arbitrary Python leaves require filtering.

Abstract

from arXiv · show

JAX and PyTorch are two popular Python autodifferentiation frameworks. JAX is based around pure functions and functional programming. PyTorch has popularised the use of an object-oriented (OO) class-based syntax for defining parameterised functions, such as neural networks. That this seems like a fundamental difference means current libraries for building parameterised functions in JAX have either rejected the OO approach entirely (Stax) or have introduced OO-to-functional transformations, multiple new abstractions, and been limited in the extent to which they integrate with JAX (Flax, Haiku, Objax). Either way this OO/functional difference has been a source of tension. Here, we introduce `Equinox', a small neural network library showing how a PyTorch-like class-based approach may be admitted without sacrificing JAX-like functional programming. We provide two main ideas. One: parameterised functions are themselves represented as `PyTrees', which means that the parameterisation of a function is transparent to the JAX framework. Two: we filter a PyTree to isolate just those components that should be treated when transforming (`jit', `grad' or `vmap'-ing) a higher-order function of a parameterised function -- such as a loss function applied to a model. Overall Equinox resolves the above tension without introducing any new programmatic abstractions: only PyTrees and transformations, just as with regular JAX. Equinox is available at \url{https://github.com/patrick-kidger/equinox}.

1 Introduction

JAX combines PyTrees for representing data with higher-order transformations that manipulate pure functions. Equinox applies these ideas to provide class-based model syntax while preserving JAX-style functional programming.

  • JAX foundations: JAX uses PyTrees to represent arbitrarily nested data and transformations such as jit, grad, and vmap to manipulate pure functions.PyTrees contain nested node types and leaves; transformations act on pure functions.
  • JAX foundations: The jit-grad-vmap pattern batches a neural-network forward pass, differentiates it with respect to model parameters, and JIT-compiles the operation.The transformations are composable around a loss function applied to a model.
  • Object-oriented model building: Class-based parameterised functions encapsulate state and define methods parameterised by that state, offering elegant model-building syntax.The attraction is syntax rather than mutation, since out-of-place updates generally suffice.
  • Contributions: Equinox demonstrates that PyTorch-like class-based syntax can coexist with JAX-like functional programming.Its approach targets model-building while retaining functional principles.
  • Contributions: Equinox avoids new abstractions, OO-to-functional translation, integration limitations, and required library-specific jit, grad, or vmap wrappers.The authors connect this reduced complexity with easier learning and interoperability with other JAX libraries.
  • Contributions: Equinox represents parameterised functions as data and filters PyTrees to select components for transformations such as jit, grad, and vmap.These are identified as Equinox’s two main ideas.

2 Related work

Prior JAX libraries either separate parameters from pure functions or add wrappers and abstractions around class-based syntax. Equinox instead treats parameterised functions as ordinary PyTrees and filters them at transformation boundaries.

  • Init/apply approaches: Stax uses an init/apply design that separates a parameter PyTree from a pure function mapping parameters and inputs to outputs.It does not represent parameterised functions themselves.
  • Init/apply wrappers: Haiku translates OO models into init/apply functions and requires multiple transform variants, while native JAX operations may need wrapped counterparts.The cited variants include transform_with_state, multi_transform, and multi_transform_with_state.
  • Init/apply wrappers: Flax combines class-based syntax with init/apply boundaries, custom parameter groups, wrapped JAX transformations, and wrapped operations such as scan.These features introduce substantial complexity around JAX API boundaries.
  • Equinox’s position: Equinox’s filter_jit and filter_grad are convenience operations over arbitrary PyTrees rather than transformations coupled to a special parameter representation.The distinction preserves compatibility with arbitrary JAX code.
  • Complete wrappers: Objax uses a wrapped API and abstractions such as variables and modules, limiting compatibility with third-party JAX libraries.Its framework-oriented approach largely forgoes interaction with native JAX.
  • Related systems: Equinox is philosophically similar to jax.tree_util.Partial but adds convenient class-based syntax for readability and composability.Flux.jl and Swift for TensorFlow are also cited as related class-based PyTree-like systems.

3 Parameterised functions as data

Equinox represents parameterised functions as immutable class instances that are also PyTrees, making their state transparent to JAX transformations. This callable-PyTree design lets higher-order functions differentiate and compile parameterised functions directly.

  • Equinox class instances are immutable and registered as custom PyTree nodes so JAX can serialise and transparently process their state.Parameters are updated out-of-place.
  • Equinox represents each parameterised function as an instance of a class, with separate classes defining different function families.
  • Class state represents the parameterisation θ, while methods define the forward operation as a pure parameterised function.Unbound methods define (θ, x) 7→fθ(x); bound methods define fθ.
  • Differentiating a higher-order function with respect to fθ returns a corresponding class instance whose parameter attribute contains the parameter gradient.The example applies this to differentiating a loss function with respect to a parameterised function.
  • Because the callable is pure when self is included, arbitrary JAX transformations such as jit can be used anywhere in the model’s forward pass.The paper contrasts this with previous libraries’ more restricted transformation use.
  • The callable-PyTree approach provides this functionality without the complexities and abstractions associated with earlier libraries.The paper notes that the same approach requires no OO-to-functional transform.

4 Filtering

Equinox filters PyTrees at transformation boundaries to separate differentiable or compiled components from static components. This supports arbitrary Python-valued parameters and selective treatment of model parameters without library-specific transformations.

  • The filtering section identifies a limitation: the simplest example assumes every PyTree leaf has JAX differentiation rules.This assumption does not hold when parameterising functions by arbitrary Python types.
  • Filtering handles parameters that JAX cannot differentiate and permits JIT compilation or differentiation with respect to only selected parameters.The motivating cases include arbitrary Python callables, differentiating floating-point arrays, and freezing some parameters.
  • Filtering partitions a model PyTree into parameters and static pieces, then reconstructs the model across a JAX API boundary.Each piece retains the model structure while containing relevant leaves and dummy values for the others.
  • Filtered transformations are wrappers for common use cases, while the underlying partitioning occurs primarily at jit, grad, and vmap API boundaries.
  • Unlike specially wrapped transformations in previous libraries, filter_grad operates on arbitrary PyTrees and is not coupled to the model.The model remains a PyTree like any other.

5 Further topics

Equinox is designed to remain simple and compatible with native JAX by treating modules and bound methods as ordinary PyTrees. It also supports explicitly returned updated state and includes a small neural-network library.

  • Equinox introduces no abstractions beyond PyTrees and transformations and aims to feel fully compatible with the main JAX library.Its feature test is whether a feature could plausibly belong in JAX itself.
  • Filtering occurs at each transformation’s call site rather than as metadata embedded in the PyTree structure.This aligns with native JAX and preserves compatibility with future JAX transformations.
  • Equinox Modules are ordinary PyTrees and are never special-cased through a separate module system.
  • All bound methods are treated as PyTrees whose single child subtree is their implicit self parameter.This extends the callable-PyTree treatment beyond __call__.
  • State updated outside gradient descent, such as batch-normalisation statistics, may be mutated and then returned from JAX API boundaries.
  • Equinox includes a small equinox.nn library as a convenience and demonstration of its parameterised-function machinery.The paper’s main focus remains representing and manipulating parameterised functions as data.

6 Conclusion

Equinox demonstrates that a PyTorch-like class-based API can build parameterised models without sacrificing JAX-style functional programming.

  • Equinox combines a PyTorch-like class-based model API with JAX-style functional programming.
  • The demonstrated compatibility applies to building models understood as parameterised functions.
  • The paper’s conclusion presents this combination as the central demonstration of Equinox.

A Example

The example defines a custom parameterised Equinox function by composing linear layers and adding a bias, then applies filtered JAX transformations to its loss. It illustrates how model PyTree leaves can include nondifferentiable or non-JIT-able components.

  • Module definition: A custom parameterised function is defined as an Equinox module with explicitly declared attributes.The example uses an eqx.Module class and specifies the module’s attributes.
  • Module definition: The module composes two linear layers, an arbitrary callable activation, and a directly defined bias parameter.The layers are stored in a nested list, the activation is a callable object, and the bias is a JAX array.
  • Initialization: The example uses JAX modules for neural-network operations, arrays, randomness, and initialization before evaluating the loss.It imports Equinox, JAX neural-network and NumPy modules, and JAX random utilities; keys initialize data and the model.
  • Forward pass: The forward pass applies the activation after each non-final layer and adds the bias to the final layer output.The implementation iterates over all but the last layer, then returns the final layer result plus the bias.
  • Transformations and loss: Filtered JIT and gradient transformations are applied to a loss function that vmaps the model over inputs and computes mean squared error.The loss evaluates vectorized predictions and returns the mean of squared prediction errors.
Loading 2111.00254v1…