Texture Neural Cellular Automata ¶

Installation¶

You will need Python 3.12 or later, and a working JAX installation. For example, you can install JAX with:

In [1]:

%pip install -U "jax[cuda]"

%pip install -U "jax[cuda]"

/home/faldor_google_com/dev/cax/.venv/bin/python3: No module named pip

Note: you may need to restart the kernel to use updated packages.

Then, install CAX from PyPi:

In [2]:

%pip install -U "cax[examples]"

%pip install -U "cax[examples]"

/home/faldor_google_com/dev/cax/.venv/bin/python3: No module named pip

Note: you may need to restart the kernel to use updated packages.

Import¶

In [3]:

import io
import random
import time
from functools import partial

import jax
import jax.numpy as jnp
import mediapy
import numpy as np
import optax
import PIL.Image
import requests
import torch
import torchvision.models as models
from flax import nnx
from jax import Array

from cax.core import ComplexSystem
from cax.core.perceive import ConvPerceive, grad2_kernel, grad_kernel, identity_kernel
from cax.core.update import ResidualUpdate
from cax.nn.pool import Pool
from cax.utils import clip_and_uint8

import io import random import time from functools import partial import jax import jax.numpy as jnp import mediapy import numpy as np import optax import PIL.Image import requests import torch import torchvision.models as models from flax import nnx from jax import Array from cax.core import ComplexSystem from cax.core.perceive import ConvPerceive, grad2_kernel, grad_kernel, identity_kernel from cax.core.update import ResidualUpdate from cax.nn.pool import Pool from cax.utils import clip_and_uint8

Configuration¶

In [4]:

seed = 0

channel_size = 12
num_kernels = 4
hidden_size = 96
cell_dropout_rate = 0.5

step_choices = [32, 48, 64, 80, 96]
num_steps = 256
pool_size = 256
batch_size = 4
learning_rate = 1e-3

spatial_dims = (128, 128)
overflow_weight = 1.0
num_projections = 32

key = jax.random.key(seed)
rngs = nnx.Rngs(seed)

seed = 0 channel_size = 12 num_kernels = 4 hidden_size = 96 cell_dropout_rate = 0.5 step_choices = [32, 48, 64, 80, 96] num_steps = 256 pool_size = 256 batch_size = 4 learning_rate = 1e-3 spatial_dims = (128, 128) overflow_weight = 1.0 num_projections = 32 key = jax.random.key(seed) rngs = nnx.Rngs(seed)

Target texture¶

In [5]:

url = "https://www.robots.ox.ac.uk/~vgg/data/dtd/thumbs/dotted/dotted_0201.jpg"
response = requests.get(url)
target_pil = PIL.Image.open(io.BytesIO(response.content)).convert("RGB")
target_pil = target_pil.resize(spatial_dims, resample=PIL.Image.Resampling.LANCZOS)

target_image = jnp.array(target_pil, dtype=jnp.float32) / 255.0

mediapy.show_image(target_image)

url = "https://www.robots.ox.ac.uk/~vgg/data/dtd/thumbs/dotted/dotted_0201.jpg" response = requests.get(url) target_pil = PIL.Image.open(io.BytesIO(response.content)).convert("RGB") target_pil = target_pil.resize(spatial_dims, resample=PIL.Image.Resampling.LANCZOS) target_image = jnp.array(target_pil, dtype=jnp.float32) / 255.0 mediapy.show_image(target_image)

VGG16 style model¶

In [6]:

STYLE_LAYERS = [1, 6, 11, 18, 25]
MAX_LAYER = max(STYLE_LAYERS) + 1


class VGGFeatureExtractor(nnx.Module):
    """VGG16 feature extractor that returns activations at specified style layers."""

    def __init__(self, torch_features: torch.nn.Sequential, *, rngs: nnx.Rngs):
        """Initialize from a torchvision VGG features sequential module.

        Args:
                torch_features: The `.features` attribute of a torchvision VGG16 model.
                rngs: Flax NNX random number generators.

        """
        convs = []
        layer_types = []

        for i in range(MAX_LAYER):
            layer = torch_features[i]
            if isinstance(layer, torch.nn.Conv2d):
                conv = nnx.Conv(
                    in_features=layer.in_channels,
                    out_features=layer.out_channels,
                    kernel_size=(layer.kernel_size[0], layer.kernel_size[1]),
                    strides=(layer.stride[0], layer.stride[1]),
                    padding=(
                        (layer.padding[0], layer.padding[0]),
                        (layer.padding[1], layer.padding[1]),
                    ),
                    use_bias=layer.bias is not None,
                    rngs=rngs,
                )
                weight_np = layer.weight.detach().cpu().numpy()
                conv.kernel.value = jnp.array(np.transpose(weight_np, (2, 3, 1, 0)))
                if layer.bias is not None:
                    conv.bias.value = jnp.array(layer.bias.detach().cpu().numpy())
                convs.append(conv)
                layer_types.append("conv")
            elif isinstance(layer, torch.nn.ReLU):
                layer_types.append("relu")
            elif isinstance(layer, torch.nn.MaxPool2d):
                layer_types.append("maxpool")
            else:
                raise ValueError(f"Unexpected layer type: {type(layer)}")

        self.convs = nnx.List(convs)
        self.layer_types = layer_types

    def __call__(self, x: jax.Array) -> list[jax.Array]:
        """Extract features at style layers.

        Args:
                x: Input images with shape (batch, height, width, 3) in [0, 1] range.

        Returns:
                List of feature maps at each style layer index.

        """
        mean = jnp.array([0.485, 0.456, 0.406])
        std = jnp.array([0.229, 0.224, 0.225])
        x = (x - mean) / std

        features = []
        conv_idx = 0
        for i, layer_type in enumerate(self.layer_types):
            if layer_type == "conv":
                x = self.convs[conv_idx](x)
                conv_idx += 1
            elif layer_type == "relu":
                x = nnx.relu(x)
            elif layer_type == "maxpool":
                x = nnx.max_pool(x, window_shape=(2, 2), strides=(2, 2))

            if i in STYLE_LAYERS:
                features.append(x)

        return features


vgg16_torch = models.vgg16(weights="IMAGENET1K_V1").features.eval()
vgg = VGGFeatureExtractor(vgg16_torch, rngs=rngs)
del vgg16_torch

STYLE_LAYERS = [1, 6, 11, 18, 25] MAX_LAYER = max(STYLE_LAYERS) + 1 class VGGFeatureExtractor(nnx.Module): """VGG16 feature extractor that returns activations at specified style layers.""" def __init__(self, torch_features: torch.nn.Sequential, *, rngs: nnx.Rngs): """Initialize from a torchvision VGG features sequential module. Args: torch_features: The `.features` attribute of a torchvision VGG16 model. rngs: Flax NNX random number generators. """ convs = [] layer_types = [] for i in range(MAX_LAYER): layer = torch_features[i] if isinstance(layer, torch.nn.Conv2d): conv = nnx.Conv( in_features=layer.in_channels, out_features=layer.out_channels, kernel_size=(layer.kernel_size[0], layer.kernel_size[1]), strides=(layer.stride[0], layer.stride[1]), padding=( (layer.padding[0], layer.padding[0]), (layer.padding[1], layer.padding[1]), ), use_bias=layer.bias is not None, rngs=rngs, ) weight_np = layer.weight.detach().cpu().numpy() conv.kernel.value = jnp.array(np.transpose(weight_np, (2, 3, 1, 0))) if layer.bias is not None: conv.bias.value = jnp.array(layer.bias.detach().cpu().numpy()) convs.append(conv) layer_types.append("conv") elif isinstance(layer, torch.nn.ReLU): layer_types.append("relu") elif isinstance(layer, torch.nn.MaxPool2d): layer_types.append("maxpool") else: raise ValueError(f"Unexpected layer type: {type(layer)}") self.convs = nnx.List(convs) self.layer_types = layer_types def __call__(self, x: jax.Array) -> list[jax.Array]: """Extract features at style layers. Args: x: Input images with shape (batch, height, width, 3) in [0, 1] range. Returns: List of feature maps at each style layer index. """ mean = jnp.array([0.485, 0.456, 0.406]) std = jnp.array([0.229, 0.224, 0.225]) x = (x - mean) / std features = [] conv_idx = 0 for i, layer_type in enumerate(self.layer_types): if layer_type == "conv": x = self.convs[conv_idx](x) conv_idx += 1 elif layer_type == "relu": x = nnx.relu(x) elif layer_type == "maxpool": x = nnx.max_pool(x, window_shape=(2, 2), strides=(2, 2)) if i in STYLE_LAYERS: features.append(x) return features vgg16_torch = models.vgg16(weights="IMAGENET1K_V1").features.eval() vgg = VGGFeatureExtractor(vgg16_torch, rngs=rngs) del vgg16_torch

Style loss¶

In [7]:

def sorted_projections(features: jax.Array, directions: jax.Array) -> jax.Array:
    """Project every position's channel vector onto random directions and sort each one.

    Sorting is what makes the comparison a texture statistic: it discards where each
    feature vector sits and keeps only the distribution of values.
    """
    b, h, w, c = features.shape
    projected = jnp.einsum("bnc,cp->bpn", features.reshape(b, h * w, c), directions)
    return jnp.sort(projected, axis=-1)


def style_loss(
    source_features: list[jax.Array], target_features: list[jax.Array], key: jax.Array
) -> jax.Array:
    """Sliced optimal transport between the source and target clouds of feature vectors.

    Comparing two high-dimensional distributions directly is expensive, so each is
    pushed onto a handful of random one-dimensional directions, where matching sorted
    samples is exactly the optimal transport cost.
    """
    loss = jnp.array(0.0)
    keys = jax.random.split(key, len(source_features))
    for source, target, subkey in zip(
        source_features, target_features, keys, strict=True
    ):
        directions = jax.random.normal(subkey, (source.shape[-1], num_projections))
        directions /= jnp.linalg.norm(directions, axis=0, keepdims=True)
        loss += jnp.sum(
            jnp.square(
                sorted_projections(source, directions)
                - sorted_projections(target, directions)
            )
        )
    return loss


target_features = vgg(target_image[None])

def sorted_projections(features: jax.Array, directions: jax.Array) -> jax.Array: """Project every position's channel vector onto random directions and sort each one. Sorting is what makes the comparison a texture statistic: it discards where each feature vector sits and keeps only the distribution of values. """ b, h, w, c = features.shape projected = jnp.einsum("bnc,cp->bpn", features.reshape(b, h * w, c), directions) return jnp.sort(projected, axis=-1) def style_loss( source_features: list[jax.Array], target_features: list[jax.Array], key: jax.Array ) -> jax.Array: """Sliced optimal transport between the source and target clouds of feature vectors. Comparing two high-dimensional distributions directly is expensive, so each is pushed onto a handful of random one-dimensional directions, where matching sorted samples is exactly the optimal transport cost. """ loss = jnp.array(0.0) keys = jax.random.split(key, len(source_features)) for source, target, subkey in zip( source_features, target_features, keys, strict=True ): directions = jax.random.normal(subkey, (source.shape[-1], num_projections)) directions /= jnp.linalg.norm(directions, axis=0, keepdims=True) loss += jnp.sum( jnp.square( sorted_projections(source, directions) - sorted_projections(target, directions) ) ) return loss target_features = vgg(target_image[None])

Instantiate system¶

In [8]:

class TextureNCA(ComplexSystem):
    """Texture Neural Cellular Automata class."""

    def __init__(self, *, rngs: nnx.Rngs):
        """Initialize Texture NCA.

        Args:
                rngs: rng key.

        """
        self.perceive = ConvPerceive(
            channel_size=channel_size,
            perception_size=num_kernels * channel_size,
            feature_group_count=channel_size,
            padding="CIRCULAR",
            rngs=rngs,
        )
        self.update = ResidualUpdate(
            num_spatial_dims=2,
            channel_size=channel_size,
            perception_size=num_kernels * channel_size,
            hidden_layer_sizes=(hidden_size,),
            cell_dropout_rate=cell_dropout_rate,
            zeros_init=True,
            rngs=rngs,
        )

        # Initialize kernel: identity + sobel x/y + laplacian
        kernel = jnp.concatenate(
            [
                identity_kernel(num_dims=2),
                grad_kernel(num_dims=2),
                grad2_kernel(num_dims=2),
            ],
            axis=-1,
        )
        kernel = jnp.expand_dims(
            jnp.concatenate([kernel] * channel_size, axis=-1), axis=-2
        )
        self.perceive.conv.kernel[...] = kernel

    def _step(self, state: Array, input: Array | None = None) -> Array:
        perception = self.perceive(state)
        next_state = self.update(state, perception, input)

        return next_state

    @nnx.jit
    def render(self, state):
        """Render state to RGB."""
        rgb = state[..., :3] + 0.5
        return clip_and_uint8(rgb)

class TextureNCA(ComplexSystem): """Texture Neural Cellular Automata class.""" def __init__(self, *, rngs: nnx.Rngs): """Initialize Texture NCA. Args: rngs: rng key. """ self.perceive = ConvPerceive( channel_size=channel_size, perception_size=num_kernels * channel_size, feature_group_count=channel_size, padding="CIRCULAR", rngs=rngs, ) self.update = ResidualUpdate( num_spatial_dims=2, channel_size=channel_size, perception_size=num_kernels * channel_size, hidden_layer_sizes=(hidden_size,), cell_dropout_rate=cell_dropout_rate, zeros_init=True, rngs=rngs, ) # Initialize kernel: identity + sobel x/y + laplacian kernel = jnp.concatenate( [ identity_kernel(num_dims=2), grad_kernel(num_dims=2), grad2_kernel(num_dims=2), ], axis=-1, ) kernel = jnp.expand_dims( jnp.concatenate([kernel] * channel_size, axis=-1), axis=-2 ) self.perceive.conv.kernel[...] = kernel def _step(self, state: Array, input: Array | None = None) -> Array: perception = self.perceive(state) next_state = self.update(state, perception, input) return next_state @nnx.jit def render(self, state): """Render state to RGB.""" rgb = state[..., :3] + 0.5 return clip_and_uint8(rgb)

In [9]:

cs = TextureNCA(rngs=rngs)

cs = TextureNCA(rngs=rngs)

In [10]:

params = nnx.state(cs, nnx.Param)
print("Number of params:", sum(x.size for x in jax.tree.leaves(params)))

params = nnx.state(cs, nnx.Param) print("Number of params:", sum(x.size for x in jax.tree.leaves(params)))

Number of params: 6300

Sample initial state¶

In [11]:

def sample_state():
    """Sample an initial state (all zeros)."""
    return jnp.zeros(spatial_dims + (channel_size,))

def sample_state(): """Sample an initial state (all zeros).""" return jnp.zeros(spatial_dims + (channel_size,))

Train¶

Pool¶

In [12]:

state = jax.vmap(lambda _: sample_state())(jnp.zeros(pool_size))

pool = Pool.create({"state": state})

state = jax.vmap(lambda _: sample_state())(jnp.zeros(pool_size)) pool = Pool.create({"state": state})

Optimizer¶

In [13]:

lr_sched = optax.piecewise_constant_schedule(
    init_value=learning_rate, boundaries_and_scales={1000: 0.3, 2000: 0.3}
)

optimizer = nnx.Optimizer(cs, optax.adam(learning_rate=lr_sched), wrt=nnx.Param)

lr_sched = optax.piecewise_constant_schedule( init_value=learning_rate, boundaries_and_scales={1000: 0.3, 2000: 0.3} ) optimizer = nnx.Optimizer(cs, optax.adam(learning_rate=lr_sched), wrt=nnx.Param)

Loss¶

In [14]:

def to_rgb(state: jax.Array) -> jax.Array:
    """Convert NCA state to RGB in [0, 1]."""
    return state[..., :3] + 0.5


def overflow_loss(state: jax.Array) -> jax.Array:
    """Penalize values outside [-1, 1]."""
    return jnp.abs(state - jnp.clip(state, -1.0, 1.0)).sum()


def loss_fn(cs, state, num_steps, key):
    """Loss function."""
    state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None})
    state = nnx.split_rngs(splits=batch_size)(
        nnx.vmap(
            lambda cs, state: cs(state, num_steps=num_steps),
            in_axes=(state_axes, 0),
        )
    )(cs, state)

    # Style loss on RGB output
    rgb = to_rgb(state)
    source_features = vgg(rgb)
    loss = style_loss(source_features, target_features, key)

    # Overflow regularization
    loss += overflow_weight * overflow_loss(state)

    return loss, state

def to_rgb(state: jax.Array) -> jax.Array: """Convert NCA state to RGB in [0, 1].""" return state[..., :3] + 0.5 def overflow_loss(state: jax.Array) -> jax.Array: """Penalize values outside [-1, 1].""" return jnp.abs(state - jnp.clip(state, -1.0, 1.0)).sum() def loss_fn(cs, state, num_steps, key): """Loss function.""" state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None}) state = nnx.split_rngs(splits=batch_size)( nnx.vmap( lambda cs, state: cs(state, num_steps=num_steps), in_axes=(state_axes, 0), ) )(cs, state) # Style loss on RGB output rgb = to_rgb(state) source_features = vgg(rgb) loss = style_loss(source_features, target_features, key) # Overflow regularization loss += overflow_weight * overflow_loss(state) return loss, state

Train step¶

In [15]:

@partial(nnx.jit, static_argnames=("num_steps",))
def train_step(cs, optimizer, pool, key, *, num_steps):
    """Train step."""
    sample_key, loss_key = jax.random.split(key)

    # Sample from pool
    pool_idx, batch = pool.sample(sample_key, batch_size=batch_size, replace=False)
    current_state = batch["state"]

    # Reset one sample to initial state
    new_state = sample_state()
    current_state = current_state.at[0].set(new_state)

    (loss, current_state), grad = nnx.value_and_grad(
        loss_fn, has_aux=True, argnums=nnx.DiffState(0, nnx.Param)
    )(cs, current_state, num_steps, loss_key)

    # Normalize gradients
    grad = jax.tree.map(lambda g: g / (jnp.linalg.norm(g) + 1e-8), grad)

    optimizer.update(cs, grad)

    pool = pool.update(pool_idx, {"state": current_state})
    return loss, pool

@partial(nnx.jit, static_argnames=("num_steps",)) def train_step(cs, optimizer, pool, key, *, num_steps): """Train step.""" sample_key, loss_key = jax.random.split(key) # Sample from pool pool_idx, batch = pool.sample(sample_key, batch_size=batch_size, replace=False) current_state = batch["state"] # Reset one sample to initial state new_state = sample_state() current_state = current_state.at[0].set(new_state) (loss, current_state), grad = nnx.value_and_grad( loss_fn, has_aux=True, argnums=nnx.DiffState(0, nnx.Param) )(cs, current_state, num_steps, loss_key) # Normalize gradients grad = jax.tree.map(lambda g: g / (jnp.linalg.norm(g) + 1e-8), grad) optimizer.update(cs, grad) pool = pool.update(pool_idx, {"state": current_state}) return loss, pool

Main loop¶

In [16]:

num_train_steps = 5_000
print_interval = 100

losses = []
start = time.perf_counter()
for i in range(num_train_steps):
    key, subkey = jax.random.split(key)
    loss, pool = train_step(
        cs, optimizer, pool, subkey, num_steps=random.choice(step_choices)
    )
    losses.append(loss)

    if i % print_interval == 0 or i == num_train_steps - 1:
        avg_loss = sum(losses[-print_interval:]) / len(losses[-print_interval:])
        elapsed = time.perf_counter() - start
        print(f"Step {i:>4}/{num_train_steps} | {elapsed:6.1f}s | Loss {avg_loss:.3e}")

print(f"✨ Trained for {num_train_steps} steps in {time.perf_counter() - start:.0f}s")

num_train_steps = 5_000 print_interval = 100 losses = [] start = time.perf_counter() for i in range(num_train_steps): key, subkey = jax.random.split(key) loss, pool = train_step( cs, optimizer, pool, subkey, num_steps=random.choice(step_choices) ) losses.append(loss) if i % print_interval == 0 or i == num_train_steps - 1: avg_loss = sum(losses[-print_interval:]) / len(losses[-print_interval:]) elapsed = time.perf_counter() - start print(f"Step {i:>4}/{num_train_steps} | {elapsed:6.1f}s | Loss {avg_loss:.3e}") print(f"✨ Trained for {num_train_steps} steps in {time.perf_counter() - start:.0f}s")

W0904 21:45:57.121569  486398 cuda_timer.cc:88] Delay kernel timed out: measured time has sub-optimal accuracy. There may be a missing warmup execution, please investigate in Nsight Systems.

Step    0/5000 |    6.8s | Loss 7.789e+06

Step  100/5000 |   23.9s | Loss 6.524e+07

Step  200/5000 |   27.9s | Loss 6.205e+06

Step  300/5000 |   31.9s | Loss 2.377e+06

Step  400/5000 |   35.8s | Loss 1.657e+06

Step  500/5000 |   39.7s | Loss 1.120e+06

Step  600/5000 |   43.7s | Loss 8.448e+05

Step  700/5000 |   47.6s | Loss 7.356e+05

Step  800/5000 |   51.6s | Loss 7.163e+05

Step  900/5000 |   55.5s | Loss 6.151e+05

Step 1000/5000 |   59.6s | Loss 5.017e+05

Step 1100/5000 |   63.8s | Loss 3.660e+05

Step 1200/5000 |   67.9s | Loss 3.216e+05

Step 1300/5000 |   72.0s | Loss 3.048e+05

Step 1400/5000 |   76.1s | Loss 2.993e+05

Step 1500/5000 |   80.1s | Loss 2.903e+05

Step 1600/5000 |   84.0s | Loss 2.797e+05

Step 1700/5000 |   88.1s | Loss 2.621e+05

Step 1800/5000 |   92.1s | Loss 2.777e+05

Step 1900/5000 |   96.1s | Loss 2.630e+05

Step 2000/5000 |  100.6s | Loss 2.485e+05

Step 2100/5000 |  104.7s | Loss 2.207e+05

Step 2200/5000 |  108.8s | Loss 2.170e+05

Step 2300/5000 |  112.9s | Loss 2.135e+05

Step 2400/5000 |  116.9s | Loss 2.147e+05

Step 2500/5000 |  120.9s | Loss 2.083e+05

Step 2600/5000 |  124.8s | Loss 2.131e+05

Step 2700/5000 |  128.9s | Loss 2.127e+05

Step 2800/5000 |  132.8s | Loss 2.130e+05

Step 2900/5000 |  136.8s | Loss 2.011e+05

Step 3000/5000 |  140.9s | Loss 2.113e+05

Step 3100/5000 |  144.9s | Loss 2.081e+05

Step 3200/5000 |  149.1s | Loss 1.951e+05

Step 3300/5000 |  152.9s | Loss 2.005e+05

Step 3400/5000 |  156.9s | Loss 1.935e+05

Step 3500/5000 |  160.9s | Loss 1.921e+05

Step 3600/5000 |  164.8s | Loss 2.059e+05

Step 3700/5000 |  168.8s | Loss 2.027e+05

Step 3800/5000 |  172.7s | Loss 2.008e+05

Step 3900/5000 |  176.7s | Loss 1.871e+05

Step 4000/5000 |  180.8s | Loss 1.925e+05

Step 4100/5000 |  184.9s | Loss 1.877e+05

Step 4200/5000 |  189.0s | Loss 1.869e+05

Step 4300/5000 |  193.0s | Loss 1.977e+05

Step 4400/5000 |  197.2s | Loss 1.804e+05

Step 4500/5000 |  201.3s | Loss 1.805e+05

Step 4600/5000 |  205.4s | Loss 1.825e+05

Step 4700/5000 |  209.4s | Loss 1.905e+05

Step 4800/5000 |  213.4s | Loss 1.812e+05

Step 4900/5000 |  217.0s | Loss 1.880e+05

Step 4999/5000 |  220.9s | Loss 1.798e+05
✨ Trained for 5000 steps in 221s

Run¶

In [17]:

num_examples = 4

state_init = jax.vmap(lambda _: sample_state())(jnp.zeros(num_examples))

state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None})
state_final, states = nnx.split_rngs(splits=num_examples)(
    nnx.vmap(
        lambda cs, state_init: cs(state_init, num_steps=num_steps, return_states=True),
        in_axes=(state_axes, 0),
    )
)(cs, state_init)

num_examples = 4 state_init = jax.vmap(lambda _: sample_state())(jnp.zeros(num_examples)) state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None}) state_final, states = nnx.split_rngs(splits=num_examples)( nnx.vmap( lambda cs, state_init: cs(state_init, num_steps=num_steps, return_states=True), in_axes=(state_axes, 0), ) )(cs, state_init)

Visualize¶

In [18]:

frames_final = nnx.vmap(
    lambda cs, state: cs.render(state),
    in_axes=(None, 0),
)(cs, state_final)

mediapy.show_images(frames_final)

frames_final = nnx.vmap( lambda cs, state: cs.render(state), in_axes=(None, 0), )(cs, state_final) mediapy.show_images(frames_final)

In [19]:

states = jnp.concatenate([state_init[:, None], states], axis=1)
frames = nnx.vmap(
    lambda cs, state: cs.render(state),
    in_axes=(None, 0),
)(cs, states)

mediapy.show_videos(frames)

states = jnp.concatenate([state_init[:, None], states], axis=1) frames = nnx.vmap( lambda cs, state: cs.render(state), in_axes=(None, 0), )(cs, states) mediapy.show_videos(frames)