Neural Network Utils

cax.nn.pool

Pool module.

Pool

Bases: Pytree

A container for PyTree arrays supporting in-place updates and random sampling.

The pool holds a PyTree of arrays whose first dimension is the pool size. It can be created from a PyTree with leading batch dimension. Sampling returns indices and the sliced batch for the same indices across all leaves.

Attributes:

Name	Type	Description
size	int	Number of items in the pool (inferred from the leading dimension of the data).
data	Any	PyTree of arrays stacked along the leading dimension.

Source code in src/cax/nn/pool.py

@nnx.dataclass
class Pool(nnx.Pytree):
    """A container for PyTree arrays supporting in-place updates and random sampling.

    The pool holds a PyTree of arrays whose first dimension is the pool size. It can be
    created from a PyTree with leading batch dimension. Sampling returns indices and the
    sliced batch for the same indices across all leaves.

    Attributes:
        size: Number of items in the pool (inferred from the leading dimension of the
            data).
        data: PyTree of arrays stacked along the leading dimension.

    """

    size: int = nnx.static()
    data: Any = nnx.data()

    @classmethod
    def create(cls, data: Any) -> Self:
        """Create a new Pool instance.

        Args:
            data: PyTree whose leaves are arrays with shape `(N, ...)`, where `N` is the
                pool size.

        Returns:
            A new Pool instance with `size == N` and `data == data`.

        """
        size = jax.tree.leaves(data)[0].shape[0]
        return cls(size=size, data=data)

    @nnx.jit
    def update(self, idxs: Array, batch: Any) -> Self:
        """Update batch in the pool at the specified indices.

        Args:
            idxs: Integer indices with shape `(B,)` indicating rows to overwrite.
            batch: PyTree matching `data` leaves sliced to `(B, ...)`.

        Returns:
            Pool with the updated batch applied at `idxs` across all leaves.

        """
        self.data = jax.tree.map(
            lambda data_leaf, batch_leaf: data_leaf.at[idxs].set(batch_leaf),
            self.data,
            batch,
        )
        return self

    @nnx.jit(static_argnames=("batch_size", "replace"))
    def sample(
        self, key: Array, *, batch_size: int, replace: bool = True
    ) -> tuple[Array, Any]:
        """Sample a batch from the pool.

        With `replace=False` the indices are distinct: a later `update(idxs, batch)`
        writes each row exactly once, with no collision between duplicate indices
        deciding which write survives — at the cost of capping `batch_size` at the
        pool size.

        Args:
            key: JAX PRNG key.
            batch_size: Number of rows to sample. With `replace=False`, must not exceed
                the pool size.
            replace: Whether to sample with replacement.

        Returns:
            A tuple `(idxs, batch)` where `idxs` has shape `(batch_size,)` and `batch`
            is a PyTree with each leaf shaped `(batch_size, ...)`.

        """
        if not replace and batch_size > self.size:
            raise ValueError(
                f"batch_size ({batch_size}) must not exceed pool size ({self.size})"
            )
        idxs = jax.random.choice(key, self.size, shape=(batch_size,), replace=replace)
        batch = jax.tree.map(lambda leaf: leaf[idxs], self.data)
        return idxs, batch

create(data) classmethod

Create a new Pool instance.

Parameters:

Name	Type	Description	Default
data	Any	PyTree whose leaves are arrays with shape (N, ...), where N is the pool size.	required

Returns:

Type	Description
Self	A new Pool instance with size == N and data == data.

Source code in src/cax/nn/pool.py

@classmethod
def create(cls, data: Any) -> Self:
    """Create a new Pool instance.

    Args:
        data: PyTree whose leaves are arrays with shape `(N, ...)`, where `N` is the
            pool size.

    Returns:
        A new Pool instance with `size == N` and `data == data`.

    """
    size = jax.tree.leaves(data)[0].shape[0]
    return cls(size=size, data=data)

update(idxs, batch)

Update batch in the pool at the specified indices.

Parameters:

Name	Type	Description	Default
idxs	Array	Integer indices with shape (B,) indicating rows to overwrite.	required
batch	Any	PyTree matching data leaves sliced to (B, ...).	required

Returns:

Type	Description
Self	Pool with the updated batch applied at idxs across all leaves.

Source code in src/cax/nn/pool.py

@nnx.jit
def update(self, idxs: Array, batch: Any) -> Self:
    """Update batch in the pool at the specified indices.

    Args:
        idxs: Integer indices with shape `(B,)` indicating rows to overwrite.
        batch: PyTree matching `data` leaves sliced to `(B, ...)`.

    Returns:
        Pool with the updated batch applied at `idxs` across all leaves.

    """
    self.data = jax.tree.map(
        lambda data_leaf, batch_leaf: data_leaf.at[idxs].set(batch_leaf),
        self.data,
        batch,
    )
    return self

sample(key, *, batch_size, replace=True)

Sample a batch from the pool.

With replace=False the indices are distinct: a later update(idxs, batch) writes each row exactly once, with no collision between duplicate indices deciding which write survives — at the cost of capping batch_size at the pool size.

Parameters:

Name	Type	Description	Default
key	Array	JAX PRNG key.	required
batch_size	int	Number of rows to sample. With replace=False, must not exceed the pool size.	required
replace	bool	Whether to sample with replacement.	True

Returns:

Type	Description
Array	A tuple (idxs, batch) where idxs has shape (batch_size,) and batch
Any	is a PyTree with each leaf shaped (batch_size, ...).

Source code in src/cax/nn/pool.py

@nnx.jit(static_argnames=("batch_size", "replace"))
def sample(
    self, key: Array, *, batch_size: int, replace: bool = True
) -> tuple[Array, Any]:
    """Sample a batch from the pool.

    With `replace=False` the indices are distinct: a later `update(idxs, batch)`
    writes each row exactly once, with no collision between duplicate indices
    deciding which write survives — at the cost of capping `batch_size` at the
    pool size.

    Args:
        key: JAX PRNG key.
        batch_size: Number of rows to sample. With `replace=False`, must not exceed
            the pool size.
        replace: Whether to sample with replacement.

    Returns:
        A tuple `(idxs, batch)` where `idxs` has shape `(batch_size,)` and `batch`
        is a PyTree with each leaf shaped `(batch_size, ...)`.

    """
    if not replace and batch_size > self.size:
        raise ValueError(
            f"batch_size ({batch_size}) must not exceed pool size ({self.size})"
        )
    idxs = jax.random.choice(key, self.size, shape=(batch_size,), replace=replace)
    batch = jax.tree.map(lambda leaf: leaf[idxs], self.data)
    return idxs, batch

cax.nn.buffer

Buffer module.

Buffer

Bases: Pytree

A container for PyTree arrays with circular writes and random sampling.

The buffer stores a PyTree of arrays with a fixed capacity along the leading dimension. New batches are written sequentially with wrap-around semantics. Sampling draws indices from the subset of entries that have been written at least once.

Attributes:

Name	Type	Description
size	int	Maximum number of items stored.
data	Any	PyTree of arrays with leading dimension size.
is_full	Array	Boolean mask of shape (size,) indicating which entries are initialized.
idx	Array	Current write pointer (modulo size).

Source code in src/cax/nn/buffer.py

@nnx.dataclass
class Buffer(nnx.Pytree):
    """A container for PyTree arrays with circular writes and random sampling.

    The buffer stores a PyTree of arrays with a fixed capacity along the leading
    dimension. New batches are written sequentially with wrap-around semantics. Sampling
    draws indices from the subset of entries that have been written at least once.

    Attributes:
        size: Maximum number of items stored.
        data: PyTree of arrays with leading dimension `size`.
        is_full: Boolean mask of shape `(size,)` indicating which entries are
            initialized.
        idx: Current write pointer (modulo `size`).

    """

    size: int = nnx.static()
    data: Any = nnx.data()
    is_full: Array = nnx.data()
    idx: Array = nnx.data()

    @classmethod
    def create(cls, size: int, datum: Any) -> Self:
        """Create a new Buffer instance.

        Args:
            size: Size of the buffer.
            datum: PyTree example whose leaf dtypes/shapes are used to allocate storage.

        Returns:
            A new Buffer instance with empty storage of capacity `size`.

        """
        data = jax.tree.map(jnp.empty_like, datum)
        data = jax.tree.map(
            lambda leaf: jnp.broadcast_to(leaf[None, ...], (size, *leaf.shape)), data
        )
        return cls(
            size=size,
            data=data,
            is_full=jnp.full((size,), False, dtype=jnp.bool),
            idx=jnp.array(0, dtype=jnp.int32),
        )

    @nnx.jit
    def add(self, batch: Any) -> Self:
        """Add a batch to the buffer.

        Args:
            batch: PyTree whose leaves have shape `(B, ...)`, where `B` is the batch
                size.

        Returns:
            Buffer with the batch written at consecutive indices (with wrap-around).

        """
        batch_size = jax.tree.leaves(batch)[0].shape[0]
        idxs = self.idx + jnp.arange(batch_size)
        idxs = idxs % self.size

        self.data = jax.tree.map(
            lambda data, batch: data.at[idxs].set(batch), self.data, batch
        )
        self.is_full = self.is_full.at[idxs].set(True)
        self.idx = (self.idx + batch_size) % self.size
        return self

    @nnx.jit(static_argnames=("batch_size",))
    def sample(self, key: Array, *, batch_size: int) -> Any:
        """Sample a batch from the buffer, with replacement.

        Replay-style buffers sample i.i.d. minibatches, so duplicate rows within a batch
        are intended; the buffer is read-only under sampling, so duplicates are
        harmless.

        Args:
            key: JAX PRNG key.
            batch_size: Number of rows to sample from initialized entries.

        Returns:
            A PyTree with each leaf shaped `(batch_size, ...)`, sampled from filled
            slots.

        """
        idxs = jax.random.choice(key, self.size, shape=(batch_size,), p=self.is_full)
        batch: Any = jax.tree.map(lambda leaf: leaf[idxs], self.data)
        return batch

create(size, datum) classmethod

Create a new Buffer instance.

Parameters:

Name	Type	Description	Default
size	int	Size of the buffer.	required
datum	Any	PyTree example whose leaf dtypes/shapes are used to allocate storage.	required

Returns:

Type	Description
Self	A new Buffer instance with empty storage of capacity size.

Source code in src/cax/nn/buffer.py

@classmethod
def create(cls, size: int, datum: Any) -> Self:
    """Create a new Buffer instance.

    Args:
        size: Size of the buffer.
        datum: PyTree example whose leaf dtypes/shapes are used to allocate storage.

    Returns:
        A new Buffer instance with empty storage of capacity `size`.

    """
    data = jax.tree.map(jnp.empty_like, datum)
    data = jax.tree.map(
        lambda leaf: jnp.broadcast_to(leaf[None, ...], (size, *leaf.shape)), data
    )
    return cls(
        size=size,
        data=data,
        is_full=jnp.full((size,), False, dtype=jnp.bool),
        idx=jnp.array(0, dtype=jnp.int32),
    )

add(batch)

Add a batch to the buffer.

Parameters:

Name	Type	Description	Default
batch	Any	PyTree whose leaves have shape (B, ...), where B is the batch size.	required

Returns:

Type	Description
Self	Buffer with the batch written at consecutive indices (with wrap-around).

Source code in src/cax/nn/buffer.py

@nnx.jit
def add(self, batch: Any) -> Self:
    """Add a batch to the buffer.

    Args:
        batch: PyTree whose leaves have shape `(B, ...)`, where `B` is the batch
            size.

    Returns:
        Buffer with the batch written at consecutive indices (with wrap-around).

    """
    batch_size = jax.tree.leaves(batch)[0].shape[0]
    idxs = self.idx + jnp.arange(batch_size)
    idxs = idxs % self.size

    self.data = jax.tree.map(
        lambda data, batch: data.at[idxs].set(batch), self.data, batch
    )
    self.is_full = self.is_full.at[idxs].set(True)
    self.idx = (self.idx + batch_size) % self.size
    return self

sample(key, *, batch_size)

Sample a batch from the buffer, with replacement.

Replay-style buffers sample i.i.d. minibatches, so duplicate rows within a batch are intended; the buffer is read-only under sampling, so duplicates are harmless.

Parameters:

Name	Type	Description	Default
key	Array	JAX PRNG key.	required
batch_size	int	Number of rows to sample from initialized entries.	required

Returns:

Type	Description
Any	A PyTree with each leaf shaped (batch_size, ...), sampled from filled
Any	slots.

Source code in src/cax/nn/buffer.py

@nnx.jit(static_argnames=("batch_size",))
def sample(self, key: Array, *, batch_size: int) -> Any:
    """Sample a batch from the buffer, with replacement.

    Replay-style buffers sample i.i.d. minibatches, so duplicate rows within a batch
    are intended; the buffer is read-only under sampling, so duplicates are
    harmless.

    Args:
        key: JAX PRNG key.
        batch_size: Number of rows to sample from initialized entries.

    Returns:
        A PyTree with each leaf shaped `(batch_size, ...)`, sampled from filled
        slots.

    """
    idxs = jax.random.choice(key, self.size, shape=(batch_size,), p=self.is_full)
    batch: Any = jax.tree.map(lambda leaf: leaf[idxs], self.data)
    return batch

cax.nn.vae

Variational Autoencoder module.

Encoder

Bases: Module

Encoder module for the VAE.

Applies a stack of strided convolutions followed by linear layers to produce mean and log-variance parameters of a diagonal Gaussian in latent space.

Source code in src/cax/nn/vae.py

class Encoder(nnx.Module):
    """Encoder module for the VAE.

    Applies a stack of strided convolutions followed by linear layers to produce mean
    and log-variance parameters of a diagonal Gaussian in latent space.
    """

    def __init__(
        self,
        *,
        spatial_dims: Sequence[int],
        features: Sequence[int],
        latent_size: int,
        rngs: nnx.Rngs,
    ):
        """Initialize the Encoder module.

        Args:
            spatial_dims: Spatial dimensions of the input.
            features: Sequence of feature sizes for convolutional layers.
            latent_size: Size of the latent space.
            rngs: rng key.

        """
        self.features = features
        self.latent_size = latent_size

        self.convs = nnx.List(
            [
                nnx.Conv(
                    in_features=in_features,
                    out_features=out_features,
                    kernel_size=(3, 3),
                    strides=(2, 2),
                    padding="SAME",
                    rngs=rngs,
                )
                for in_features, out_features in zip(
                    self.features[:-1], self.features[1:], strict=True
                )
            ]
        )

        flattened_size = spatial_dims[0] * spatial_dims[1] * self.features[-1]
        for _ in range(len(self.features) - 1):
            flattened_size //= 4

        self.linear = nnx.Linear(
            in_features=flattened_size, out_features=flattened_size, rngs=rngs
        )
        self.mean = nnx.Linear(
            in_features=flattened_size, out_features=self.latent_size, rngs=rngs
        )
        self.logvar = nnx.Linear(
            in_features=flattened_size, out_features=self.latent_size, rngs=rngs
        )
        self.rngs = rngs

    def __call__(self, x: Array) -> tuple[Array, Array]:
        """Forward pass of the encoder.

        Args:
            x: Input tensor with shape `(..., H, W, channel_size)`.

        Returns:
            Tuple `(mean, logvar)` each with shape `(..., latent_size)`.

        """
        for conv in self.convs:
            x = jax.nn.relu(conv(x))
        x = jnp.reshape(x, (*x.shape[:-3], -1))
        x = jax.nn.relu(self.linear(x))
        mean = self.mean(x)
        logvar = self.logvar(x)
        return mean, logvar

    def reparameterize(self, mean: Array, logvar: Array) -> Array:
        """Perform the reparameterization trick.

        Args:
            mean: Mean of the latent distribution.
            logvar: Log variance of the latent distribution.

        Returns:
            Sampled latent vector with shape matching `mean`.

        """
        return mean + jnp.exp(logvar * 0.5) * jax.random.normal(
            self.rngs(), shape=mean.shape
        )

__init__(*, spatial_dims, features, latent_size, rngs)

Initialize the Encoder module.

Parameters:

Name	Type	Description	Default
spatial_dims	Sequence[int]	Spatial dimensions of the input.	required
features	Sequence[int]	Sequence of feature sizes for convolutional layers.	required
latent_size	int	Size of the latent space.	required
rngs	Rngs	rng key.	required

Source code in src/cax/nn/vae.py

def __init__(
    self,
    *,
    spatial_dims: Sequence[int],
    features: Sequence[int],
    latent_size: int,
    rngs: nnx.Rngs,
):
    """Initialize the Encoder module.

    Args:
        spatial_dims: Spatial dimensions of the input.
        features: Sequence of feature sizes for convolutional layers.
        latent_size: Size of the latent space.
        rngs: rng key.

    """
    self.features = features
    self.latent_size = latent_size

    self.convs = nnx.List(
        [
            nnx.Conv(
                in_features=in_features,
                out_features=out_features,
                kernel_size=(3, 3),
                strides=(2, 2),
                padding="SAME",
                rngs=rngs,
            )
            for in_features, out_features in zip(
                self.features[:-1], self.features[1:], strict=True
            )
        ]
    )

    flattened_size = spatial_dims[0] * spatial_dims[1] * self.features[-1]
    for _ in range(len(self.features) - 1):
        flattened_size //= 4

    self.linear = nnx.Linear(
        in_features=flattened_size, out_features=flattened_size, rngs=rngs
    )
    self.mean = nnx.Linear(
        in_features=flattened_size, out_features=self.latent_size, rngs=rngs
    )
    self.logvar = nnx.Linear(
        in_features=flattened_size, out_features=self.latent_size, rngs=rngs
    )
    self.rngs = rngs

__call__(x)

Forward pass of the encoder.

Parameters:

Name	Type	Description	Default
x	Array	Input tensor with shape (..., H, W, channel_size).	required

Returns:

Type	Description
tuple[Array, Array]	Tuple (mean, logvar) each with shape (..., latent_size).

Source code in src/cax/nn/vae.py

def __call__(self, x: Array) -> tuple[Array, Array]:
    """Forward pass of the encoder.

    Args:
        x: Input tensor with shape `(..., H, W, channel_size)`.

    Returns:
        Tuple `(mean, logvar)` each with shape `(..., latent_size)`.

    """
    for conv in self.convs:
        x = jax.nn.relu(conv(x))
    x = jnp.reshape(x, (*x.shape[:-3], -1))
    x = jax.nn.relu(self.linear(x))
    mean = self.mean(x)
    logvar = self.logvar(x)
    return mean, logvar

reparameterize(mean, logvar)

Perform the reparameterization trick.

Parameters:

Name	Type	Description	Default
mean	Array	Mean of the latent distribution.	required
logvar	Array	Log variance of the latent distribution.	required

Returns:

Type	Description
Array	Sampled latent vector with shape matching mean.

Source code in src/cax/nn/vae.py

def reparameterize(self, mean: Array, logvar: Array) -> Array:
    """Perform the reparameterization trick.

    Args:
        mean: Mean of the latent distribution.
        logvar: Log variance of the latent distribution.

    Returns:
        Sampled latent vector with shape matching `mean`.

    """
    return mean + jnp.exp(logvar * 0.5) * jax.random.normal(
        self.rngs(), shape=mean.shape
    )

Decoder

Bases: Module

Decoder module for the VAE.

Maps latent vectors back to image space using a linear layer followed by transposed convolutions.

Source code in src/cax/nn/vae.py

class Decoder(nnx.Module):
    """Decoder module for the VAE.

    Maps latent vectors back to image space using a linear layer followed by transposed
    convolutions.
    """

    def __init__(
        self,
        *,
        spatial_dims: Sequence[int],
        features: Sequence[int],
        latent_size: int,
        rngs: nnx.Rngs,
    ):
        """Initialize the Decoder module.

        Args:
            spatial_dims: Spatial dimensions of the output.
            features: Sequence of feature sizes for transposed convolutional layers.
            latent_size: Size of the latent space.
            rngs: rng key.

        """
        self.features = features
        self.latent_size = latent_size

        self._spatial_dims = tuple(
            dim // (2 ** (len(self.features) - 1)) for dim in spatial_dims[:2]
        )

        flattened_size = (
            self._spatial_dims[0] * self._spatial_dims[1] * self.features[0]
        )

        self.linear = nnx.Linear(
            in_features=self.latent_size, out_features=flattened_size, rngs=rngs
        )

        self.convs = nnx.List(
            [
                nnx.ConvTranspose(
                    in_features=in_features,
                    out_features=out_features,
                    kernel_size=(3, 3),
                    strides=(2, 2),
                    padding="SAME",
                    rngs=rngs,
                )
                for in_features, out_features in zip(
                    self.features[:-1], self.features[1:], strict=True
                )
            ]
        )

    def __call__(self, z: Array) -> Array:
        """Forward pass of the decoder.

        Args:
            z: Latent vector with shape `(..., latent_size)`.

        Returns:
            Reconstructed output tensor with shape `(..., H, W, channel_size)`.

        """
        x = jax.nn.relu(self.linear(z))
        x = jnp.reshape(x, x.shape[:-1] + self._spatial_dims + (self.features[0],))
        for conv in self.convs[:-1]:
            x = jax.nn.relu(conv(x))
        x = self.convs[-1](x)
        return x

__init__(*, spatial_dims, features, latent_size, rngs)

Initialize the Decoder module.

Parameters:

Name	Type	Description	Default
spatial_dims	Sequence[int]	Spatial dimensions of the output.	required
features	Sequence[int]	Sequence of feature sizes for transposed convolutional layers.	required
latent_size	int	Size of the latent space.	required
rngs	Rngs	rng key.	required

Source code in src/cax/nn/vae.py

def __init__(
    self,
    *,
    spatial_dims: Sequence[int],
    features: Sequence[int],
    latent_size: int,
    rngs: nnx.Rngs,
):
    """Initialize the Decoder module.

    Args:
        spatial_dims: Spatial dimensions of the output.
        features: Sequence of feature sizes for transposed convolutional layers.
        latent_size: Size of the latent space.
        rngs: rng key.

    """
    self.features = features
    self.latent_size = latent_size

    self._spatial_dims = tuple(
        dim // (2 ** (len(self.features) - 1)) for dim in spatial_dims[:2]
    )

    flattened_size = (
        self._spatial_dims[0] * self._spatial_dims[1] * self.features[0]
    )

    self.linear = nnx.Linear(
        in_features=self.latent_size, out_features=flattened_size, rngs=rngs
    )

    self.convs = nnx.List(
        [
            nnx.ConvTranspose(
                in_features=in_features,
                out_features=out_features,
                kernel_size=(3, 3),
                strides=(2, 2),
                padding="SAME",
                rngs=rngs,
            )
            for in_features, out_features in zip(
                self.features[:-1], self.features[1:], strict=True
            )
        ]
    )

__call__(z)

Forward pass of the decoder.

Parameters:

Name	Type	Description	Default
z	Array	Latent vector with shape (..., latent_size).	required

Returns:

Type	Description
Array	Reconstructed output tensor with shape (..., H, W, channel_size).

Source code in src/cax/nn/vae.py

def __call__(self, z: Array) -> Array:
    """Forward pass of the decoder.

    Args:
        z: Latent vector with shape `(..., latent_size)`.

    Returns:
        Reconstructed output tensor with shape `(..., H, W, channel_size)`.

    """
    x = jax.nn.relu(self.linear(z))
    x = jnp.reshape(x, x.shape[:-1] + self._spatial_dims + (self.features[0],))
    for conv in self.convs[:-1]:
        x = jax.nn.relu(conv(x))
    x = self.convs[-1](x)
    return x

VAE

Bases: Module

Variational Autoencoder module.

Combines an encoder and decoder with a reparameterization sampler for training with the evidence lower bound (ELBO).

Source code in src/cax/nn/vae.py

class VAE(nnx.Module):
    """Variational Autoencoder module.

    Combines an encoder and decoder with a reparameterization sampler for training with
    the evidence lower bound (ELBO).

    """

    def __init__(
        self,
        *,
        spatial_dims: tuple[int, int],
        features: Sequence[int],
        latent_size: int,
        rngs: nnx.Rngs,
    ):
        """Initialize the VAE module.

        Args:
            spatial_dims: Spatial dimensions of the input/output.
            features: Sequence of feature sizes for encoder and decoder.
            latent_size: Size of the latent space.
            rngs: rng key.

        """
        super().__init__()
        self.encoder = Encoder(
            spatial_dims=spatial_dims,
            features=features,
            latent_size=latent_size,
            rngs=rngs,
        )
        self.decoder = Decoder(
            spatial_dims=spatial_dims,
            features=features[::-1],
            latent_size=latent_size,
            rngs=rngs,
        )

    def encode(self, x: Array) -> tuple[Array, Array, Array]:
        """Encode input to latent space.

        Args:
            x: Input tensor with shape `(..., H, W, channel_size)`.

        Returns:
            Tuple `(z, mean, logvar)` where all have shape `(..., latent_size)`.

        """
        mean, logvar = self.encoder(x)
        return self.encoder.reparameterize(mean, logvar), mean, logvar

    def decode(self, z: Array) -> Array:
        """Decode latent vector to output space.

        Args:
            z: Latent vector with shape `(..., latent_size)`.

        Returns:
            Reconstructed output tensor with shape `(..., H, W, channel_size)`.

        """
        return self.decoder(z)

    def generate(self, z: Array) -> Array:
        """Generate output from latent vector.

        Args:
            z: Latent vector with shape `(..., latent_size)`.

        Returns:
            Generated output tensor with shape `(..., H, W, channel_size)` in the range
            `[0, 1]`.

        """
        return jax.nn.sigmoid(self.decoder(z))

    def __call__(self, x: Array) -> tuple[Array, Array, Array]:
        """Forward pass of the VAE.

        Args:
            x: Input tensor with shape `(..., H, W, channel_size)`.

        Returns:
            Tuple `(logits, mean, logvar)`

        """
        z, mean, logvar = self.encode(x)
        logits = self.decode(z)
        return logits, mean, logvar

__init__(*, spatial_dims, features, latent_size, rngs)

Initialize the VAE module.

Parameters:

Name	Type	Description	Default
spatial_dims	tuple[int, int]	Spatial dimensions of the input/output.	required
features	Sequence[int]	Sequence of feature sizes for encoder and decoder.	required
latent_size	int	Size of the latent space.	required
rngs	Rngs	rng key.	required

Source code in src/cax/nn/vae.py

def __init__(
    self,
    *,
    spatial_dims: tuple[int, int],
    features: Sequence[int],
    latent_size: int,
    rngs: nnx.Rngs,
):
    """Initialize the VAE module.

    Args:
        spatial_dims: Spatial dimensions of the input/output.
        features: Sequence of feature sizes for encoder and decoder.
        latent_size: Size of the latent space.
        rngs: rng key.

    """
    super().__init__()
    self.encoder = Encoder(
        spatial_dims=spatial_dims,
        features=features,
        latent_size=latent_size,
        rngs=rngs,
    )
    self.decoder = Decoder(
        spatial_dims=spatial_dims,
        features=features[::-1],
        latent_size=latent_size,
        rngs=rngs,
    )

encode(x)

Encode input to latent space.

Parameters:

Name	Type	Description	Default
x	Array	Input tensor with shape (..., H, W, channel_size).	required

Returns:

Type	Description
tuple[Array, Array, Array]	Tuple (z, mean, logvar) where all have shape (..., latent_size).

Source code in src/cax/nn/vae.py

def encode(self, x: Array) -> tuple[Array, Array, Array]:
    """Encode input to latent space.

    Args:
        x: Input tensor with shape `(..., H, W, channel_size)`.

    Returns:
        Tuple `(z, mean, logvar)` where all have shape `(..., latent_size)`.

    """
    mean, logvar = self.encoder(x)
    return self.encoder.reparameterize(mean, logvar), mean, logvar

decode(z)

Decode latent vector to output space.

Parameters:

Name	Type	Description	Default
z	Array	Latent vector with shape (..., latent_size).	required

Returns:

Type	Description
Array	Reconstructed output tensor with shape (..., H, W, channel_size).

Source code in src/cax/nn/vae.py

def decode(self, z: Array) -> Array:
    """Decode latent vector to output space.

    Args:
        z: Latent vector with shape `(..., latent_size)`.

    Returns:
        Reconstructed output tensor with shape `(..., H, W, channel_size)`.

    """
    return self.decoder(z)

generate(z)

Generate output from latent vector.

Parameters:

Name	Type	Description	Default
z	Array	Latent vector with shape (..., latent_size).	required

Returns:

Type	Description
Array	Generated output tensor with shape (..., H, W, channel_size) in the range
Array	[0, 1].

Source code in src/cax/nn/vae.py

def generate(self, z: Array) -> Array:
    """Generate output from latent vector.

    Args:
        z: Latent vector with shape `(..., latent_size)`.

    Returns:
        Generated output tensor with shape `(..., H, W, channel_size)` in the range
        `[0, 1]`.

    """
    return jax.nn.sigmoid(self.decoder(z))

__call__(x)

Forward pass of the VAE.

Parameters:

Name	Type	Description	Default
x	Array	Input tensor with shape (..., H, W, channel_size).	required

Returns:

Type	Description
tuple[Array, Array, Array]	Tuple (logits, mean, logvar)

Source code in src/cax/nn/vae.py

def __call__(self, x: Array) -> tuple[Array, Array, Array]:
    """Forward pass of the VAE.

    Args:
        x: Input tensor with shape `(..., H, W, channel_size)`.

    Returns:
        Tuple `(logits, mean, logvar)`

    """
    z, mean, logvar = self.encode(x)
    logits = self.decode(z)
    return logits, mean, logvar

kl_divergence(mean, logvar)

Compute KL divergence between latent distribution and standard normal.

Parameters:

Name	Type	Description	Default
mean	Array	Mean of the latent distribution.	required
logvar	Array	Log variance of the latent distribution.	required

Returns:

Type	Description
Array	Scalar KL divergence value (sum over last dimension).

Source code in src/cax/nn/vae.py

@nnx.jit
def kl_divergence(mean: Array, logvar: Array) -> Array:
    """Compute KL divergence between latent distribution and standard normal.

    Args:
        mean: Mean of the latent distribution.
        logvar: Log variance of the latent distribution.

    Returns:
        Scalar KL divergence value (sum over last dimension).

    """
    return -0.5 * jnp.sum(1 + logvar - jnp.square(mean) - jnp.exp(logvar))

binary_cross_entropy_with_logits(logits, labels)

Compute binary cross-entropy loss with logits.

Parameters:

Name	Type	Description	Default
logits	Array	Predicted logits.	required
labels	Array	True labels.	required

Returns:

Type	Description
Array	Summed Binary Cross-Entropy loss over the last dimension.

Source code in src/cax/nn/vae.py

@nnx.jit
def binary_cross_entropy_with_logits(logits: Array, labels: Array) -> Array:
    """Compute binary cross-entropy loss with logits.

    Args:
        logits: Predicted logits.
        labels: True labels.

    Returns:
        Summed Binary Cross-Entropy loss over the last dimension.

    """
    logits = jax.nn.log_sigmoid(logits)
    return -jnp.sum(labels * logits + (1.0 - labels) * jnp.log(-jnp.expm1(logits)))

vae_loss(logits, labels, mean, logvar)

Compute VAE loss.

Parameters:

Name	Type	Description	Default
logits	Array	Predicted logits from the decoder.	required
labels	Array	Target labels (e.g., normalized images or one-hot vectors).	required
mean	Array	Mean of the latent distribution.	required
logvar	Array	Log variance of the latent distribution.	required

Returns:

Type	Description
Array	Total VAE loss equal to mean(BCE) + mean(KL).

Source code in src/cax/nn/vae.py

@nnx.jit
def vae_loss(logits: Array, labels: Array, mean: Array, logvar: Array) -> Array:
    """Compute VAE loss.

    Args:
        logits: Predicted logits from the decoder.
        labels: Target labels (e.g., normalized images or one-hot vectors).
        mean: Mean of the latent distribution.
        logvar: Log variance of the latent distribution.

    Returns:
        Total VAE loss equal to `mean(BCE) + mean(KL)`.

    """
    bce_loss = jnp.mean(binary_cross_entropy_with_logits(logits, labels))
    kld_loss = jnp.mean(kl_divergence(mean, logvar))
    return bce_loss + kld_loss