Skip to main content

Immutable Jax pytree modules classes with easy manipulation and serialization

Project description

Noxjax

Simple pytree module classes for Jax, strongly inspired by Equinox.

  • Referential transparency via strict immutability
  • Safe serialization including hyperparameters
  • Bound methods and function transformations are also modules
  • Auxillary information in key paths for filtered transformations

Quick Examples

Modules work similar to dataclasses, but with the added benefit of being pytrees. Making them compatible with all Jax function transformations.

import noxjax as nox

class Linear(nox.Module):
    # The __init__ method is automatically generated
    w: jax.Array
    b: jax.Array

    # additional intialization methods via classmethods
    @classmethod
    def init(cls, key, dim_in, dim):
        w = jax.random.normal(key, (dim, dim_in)) * 0.02
        b = jax.numpy.zeros((dim,))
        return cls(w=w, b=b)

    def __call__(self, x):
        return self.w @ x + self.b

key = jax.random.PRNGKey(42)
key1, key2 = jax.random.split(key)

model = fj.Sequential(
    (
        Linear.init(key1, 3, 2),
        Linear.init(key2, 2, 5),
    )
)

The model can be serialized and deserialized using fj.save and fj.load.

nox.save("model.npz", model)
model = nox.load("model.npz")

Noxjax includes wrappers of the Jax function transformations, which return callable modules.

model = nox.VMap(model)
model = nox.Jit(model)

Installation

Memmpy can be installed directly from PyPI using pip. It requires Python 3.10+ and Jax 0.4.26+.

pip install Noxjax

Design

Noxjax modules sacrifice some flexibility for the sake of a unified interface and safety. Noxjax code should alway be easy to reason about and should not contain any footguns from using python magic.

  1. Everything is immutable and
  2. module fields can be either jax arrays, other modules or json-like data.

This makes it harder to use other jax libraries in Noxjax modules. It is recommended to wrap the needed functionality in a module. Most jax libraries should be compatible with Noxjax modules, since they are simply callable pytrees.

Roadmap

  • Filtered grad transformation based on key paths
  • Pretty printing for modules
  • Rule to infer static arguments in jitted functions, possibly everything except JAX arrays

See also

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

noxjax-0.1.0.tar.gz (11.8 kB view hashes)

Uploaded Source

Built Distribution

noxjax-0.1.0-py3-none-any.whl (13.0 kB view hashes)

Uploaded Python 3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page