Pinned Loading

1. jaxonmodels jaxonmodels Public

   JAX models for ML

   Python 6

2. statedict2pytree statedict2pytree Public

   Python 44 1

3. jaxonlayers jaxonlayers Public

   Additional layers for use with Equinox and JAX

   Python 1

4. An example showcasing that you can c... An example showcasing that you can combine Equinox and Flax

   1

   from collections.abc import Callable, Sequence

   2

   from functools import partial

   3

   from typing import Any, Tuple

   4

   5

   import equinox as eqx

5. Text diffusion on tinyshakespeare us... Text diffusion on tinyshakespeare using Equinox and using discrete representations instead of embedding everything

   1

   import os

   2

   3

   import clu.metrics as clum

   4

   import equinox as eqx

   5

   import flax

6. VAE in Equinox/JAX VAE in Equinox/JAX

   1

   import clu.metrics as clum

   2

   import equinox as eqx

   3

   import jax

   4

   import jax.numpy as jnp

   5

   import matplotlib.pyplot as plt