Equinox Documentation

Equinox differs from Flax in several key architectural ways:

• Abstractions: Flax introduces several new abstractions (e.g., flax.linen.Module, flax.linen.Variable, Module.setup). Equinox avoids adding new abstractions to core JAX; everything in Equinox is always just a PyTree.
• JAX Compatibility: Flax is a DSL that is generally incompatible with non-Flax code and requires using wrapped operations like flax.linen.vmap or flax.linen.scan. Equinox is compatible with arbitrary JAX code and allows the use of native JAX operations.
• Bound Methods: Bound methods of eqx.Module are PyTrees, making them easy to pass around. In Flax, passing bound methods can lead to errors or unexpected recompilations.
• Advanced Features: Equinox provides features like runtime errors and specialized PyTree manipulation (e.g., equinox.tree_at) that are not present in Flax.

Equinox (via JAX) and Julia have different strengths:

• Autodifferentiation: JAX+Equinox uses a single, robust autodifferentiation system. Julia has multiple competing implementations for both forward-mode and reverse-mode autodifferentiation and lacks robust support for higher-order autodifferentiation.
• Use Cases: Julia is superior for general-purpose programming tasks like branch-and-bound combinatorial optimization. Equinox is optimized specifically for arrays and linear algebra (NumPy-style workloads).
• Hardware: Equinox supports running on TPUs, whereas Julia generally does not.

The equinox.internal namespace (often imported as eqxi) contains advanced, undocumented features intended for privileged downstream libraries (like Diffrax) or JAX enthusiasts.

WARNING: This namespace is 'semi-public'. APIs may change or be removed at any time without stability guarantees. Use with caution.

import equinox.internal as eqxi

When passing a module containing non-JAX objects (static elements) into JAX control flow primitives like jax.lax.scan, you must use eqx.partition and eqx.combine to separate the arrays from the static parts. This allows the arrays to be part of the scan's carry while the static parts are captured via closure.

def rollout(mlp, xs):
    # Separate arrays from static elements
    arr, static = eqx.partition(mlp, eqx.is_array)
    
    def step(carry, x):
        # Recombine static elements into the module for use
        mlp = eqx.combine(carry, static)
        val = mlp(x)
        # Partition again to prepare the next carry
        carry, _ = eqx.partition(mlp, eqx.is_array)
        return carry, [val]

    _, scan_out = jax.lax.scan(step, arr, xs)
    return scan_out

To maintain type safety and predictable return types in functions, the abstract/final pattern discourages adding new methods to a concrete subclass that were not present in the superclass.

By ensuring that a concrete class only implements what was defined in its abstract ancestors, you avoid ambiguity in type signatures. For example, if a function is typed to return ConcreteArray, you avoid the question of whether it should return a specific subclass like ConcreteArrayTwo when different types are passed as arguments.

# Avoid this pattern to keep return types predictable
class ConcreteArray(eqx.Module):
    def some_method(self):
        pass

class ConcreteArrayTwo(ConcreteArray):
    def another_method(self):
        pass

To treat an array as a non-trainable buffer (similar to PyTorch buffers), wrap the array in jax.lax.stop_gradient within your model's __call__ method.

class Model(eqx.Module):
    buffer: Array
    param: Array

    def __call__(self, x):
        return self.param * x + jax.lax.stop_gradient(self.buffer)