Growing Conditional 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 time

import jax
import jax.numpy as jnp
import mediapy
import optax
from flax import nnx
from jax import Array

from cax.core import ComplexSystem
from cax.core.perceive import ConvPerceive, grad_kernel, identity_kernel
from cax.core.update import NCAUpdate
from cax.nn.pool import Pool
from cax.utils import clip_and_uint8, get_emoji_array, rgba_to_rgb

import time import jax import jax.numpy as jnp import mediapy import optax from flax import nnx from jax import Array from cax.core import ComplexSystem from cax.core.perceive import ConvPerceive, grad_kernel, identity_kernel from cax.core.update import NCAUpdate from cax.nn.pool import Pool from cax.utils import clip_and_uint8, get_emoji_array, rgba_to_rgb

Configuration¶

In [4]:

seed = 0

channel_size = 32
num_kernels = 3
hidden_size = 256
cell_dropout_rate = 0.5

num_steps = 128
pool_size = 1_024
batch_size = 8
learning_rate = 1e-3

emojis = "🐶🐱🐭🐹🐰🦊🐻🐼"
size = 40
pad_width = 16

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

seed = 0 channel_size = 32 num_kernels = 3 hidden_size = 256 cell_dropout_rate = 0.5 num_steps = 128 pool_size = 1_024 batch_size = 8 learning_rate = 1e-3 emojis = "🐶🐱🐭🐹🐰🦊🐻🐼" size = 40 pad_width = 16 key = jax.random.key(seed) rngs = nnx.Rngs(seed)

Dataset¶

In [5]:

y = jnp.array([get_emoji_array(emoji, size, pad_width) for emoji in emojis])
z = jnp.eye(y.shape[0])

mediapy.show_images(y.repeat(2, axis=-3).repeat(2, axis=-2))

y = jnp.array([get_emoji_array(emoji, size, pad_width) for emoji in emojis]) z = jnp.eye(y.shape[0]) mediapy.show_images(y.repeat(2, axis=-3).repeat(2, axis=-2))

Instantiate system¶

In [6]:

class GrowingConditionalNCA(ComplexSystem):
    """Growing Conditional Neural Cellular Automata class."""

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

        Args:
            rngs: rng key.

        """
        self.perceive = ConvPerceive(
            channel_size=channel_size,
            perception_size=num_kernels * channel_size,
            feature_group_count=channel_size,
            rngs=rngs,
        )
        self.update = NCAUpdate(
            channel_size=channel_size,
            perception_size=y.shape[0] + num_kernels * channel_size,
            hidden_layer_sizes=(hidden_size,),
            cell_dropout_rate=cell_dropout_rate,
            zeros_init=True,
            rngs=rngs,
        )

        # Initialize kernel with sobel filters
        kernel = jnp.concatenate(
            [identity_kernel(num_dims=2), grad_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:
        # Broadcast the input vector to match the state shape
        input_shape = (*state.shape[:-1], input.shape[-1])
        input = jnp.broadcast_to(input[..., None, None, :], input_shape)

        # Step
        perception = self.perceive(state)
        next_state = self.update(state, perception, input)

        return next_state

    @nnx.jit
    def render(self, state):
        """Render state to RGB."""
        rgba = state[..., -4:]
        rgb = rgba_to_rgb(rgba)

        # Clip values to valid range and convert to uint8
        return clip_and_uint8(rgb)

    @nnx.jit
    def render_rgba(self, state):
        """Render state to RGBA."""
        rgba = state[..., -4:]

        # Clip values to valid range and convert to uint8
        return clip_and_uint8(rgba)

class GrowingConditionalNCA(ComplexSystem): """Growing Conditional Neural Cellular Automata class.""" def __init__(self, *, rngs: nnx.Rngs): """Initialize Growing Conditional NCA. Args: rngs: rng key. """ self.perceive = ConvPerceive( channel_size=channel_size, perception_size=num_kernels * channel_size, feature_group_count=channel_size, rngs=rngs, ) self.update = NCAUpdate( channel_size=channel_size, perception_size=y.shape[0] + num_kernels * channel_size, hidden_layer_sizes=(hidden_size,), cell_dropout_rate=cell_dropout_rate, zeros_init=True, rngs=rngs, ) # Initialize kernel with sobel filters kernel = jnp.concatenate( [identity_kernel(num_dims=2), grad_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: # Broadcast the input vector to match the state shape input_shape = (*state.shape[:-1], input.shape[-1]) input = jnp.broadcast_to(input[..., None, None, :], input_shape) # Step perception = self.perceive(state) next_state = self.update(state, perception, input) return next_state @nnx.jit def render(self, state): """Render state to RGB.""" rgba = state[..., -4:] rgb = rgba_to_rgb(rgba) # Clip values to valid range and convert to uint8 return clip_and_uint8(rgb) @nnx.jit def render_rgba(self, state): """Render state to RGBA.""" rgba = state[..., -4:] # Clip values to valid range and convert to uint8 return clip_and_uint8(rgba)

In [7]:

cs = GrowingConditionalNCA(rngs=rngs)

cs = GrowingConditionalNCA(rngs=rngs)

In [8]:

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: 35968

Sample initial state¶

In [9]:

def sample_state(key):
    """Sample a state with a single alive cell."""
    spatial_dims = y.shape[1:3]

    # Init state
    state = jnp.zeros(spatial_dims + (channel_size,))
    mid = tuple(size // 2 for size in spatial_dims)

    # Set the center cell to alive, with hidden channels at one
    state = state.at[mid[0], mid[1], -1].set(1.0)
    state = state.at[mid[0], mid[1], :-4].set(1.0)

    # Sample a random target y
    y_idx = jax.random.choice(key, y.shape[0])
    return state, y_idx

def sample_state(key): """Sample a state with a single alive cell.""" spatial_dims = y.shape[1:3] # Init state state = jnp.zeros(spatial_dims + (channel_size,)) mid = tuple(size // 2 for size in spatial_dims) # Set the center cell to alive, with hidden channels at one state = state.at[mid[0], mid[1], -1].set(1.0) state = state.at[mid[0], mid[1], :-4].set(1.0) # Sample a random target y y_idx = jax.random.choice(key, y.shape[0]) return state, y_idx

Train¶

Pool¶

In [10]:

key, subkey = jax.random.split(key)

keys = jax.random.split(subkey, pool_size)
state, y_idx = jax.vmap(lambda key: sample_state(key))(keys)

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

key, subkey = jax.random.split(key) keys = jax.random.split(subkey, pool_size) state, y_idx = jax.vmap(lambda key: sample_state(key))(keys) pool = Pool.create({"state": state, "y_idx": y_idx})

Optimizer¶

In [11]:

lr_sched = optax.linear_schedule(
    init_value=learning_rate, end_value=0.1 * learning_rate, transition_steps=8_192
)

optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adam(learning_rate=lr_sched),
)

update_params = nnx.All(nnx.Param, nnx.PathContains("update"))
optimizer = nnx.Optimizer(cs, optimizer, wrt=update_params)

lr_sched = optax.linear_schedule( init_value=learning_rate, end_value=0.1 * learning_rate, transition_steps=8_192 ) optimizer = optax.chain( optax.clip_by_global_norm(1.0), optax.adam(learning_rate=lr_sched), ) update_params = nnx.All(nnx.Param, nnx.PathContains("update")) optimizer = nnx.Optimizer(cs, optimizer, wrt=update_params)

Loss¶

In [12]:

def mse(state, y):
    """Mean Squared Error."""
    return jnp.mean(jnp.square(state[..., -4:] - y))

def mse(state, y): """Mean Squared Error.""" return jnp.mean(jnp.square(state[..., -4:] - y))

In [13]:

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

    # Sample a random step
    idx = jax.random.randint(key, (batch_size,), num_steps // 2, num_steps)
    state = state[jnp.arange(batch_size), idx]

    loss = mse(state, y)
    return loss, state

@nnx.jit def loss_fn(cs, state, z, y, key): """Loss function.""" state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None}) _, state = nnx.split_rngs(splits=batch_size)( nnx.vmap( lambda cs, state, z: cs(state, z, num_steps=num_steps, return_states=True), in_axes=(state_axes, 0, 0), ) )(cs, state, z) # Sample a random step idx = jax.random.randint(key, (batch_size,), num_steps // 2, num_steps) state = state[jnp.arange(batch_size), idx] loss = mse(state, y) return loss, state

Train step¶

In [14]:

@nnx.jit
def train_step(cs, optimizer, pool, key):
    """Train step."""
    sample_key, sample_state_key, loss_key = jax.random.split(key, 3)

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

    # Sort by descending loss
    sort_idx = jnp.argsort(jax.vmap(mse)(current_state, current_y), descending=True)
    pool_idx = pool_idx[sort_idx]
    current_state = current_state[sort_idx]
    current_y_idx = current_y_idx[sort_idx]

    # Sample a new state to replace the worst
    new_state, new_y_idx = sample_state(sample_state_key)
    current_state = current_state.at[0].set(new_state)
    current_y_idx = current_y_idx.at[0].set(new_y_idx)
    current_y = y[current_y_idx]
    current_z = z[current_y_idx]

    (loss, current_state), grad = nnx.value_and_grad(
        loss_fn, has_aux=True, argnums=nnx.DiffState(0, update_params)
    )(cs, current_state, current_z, current_y, loss_key)
    optimizer.update(cs, grad)

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

@nnx.jit def train_step(cs, optimizer, pool, key): """Train step.""" sample_key, sample_state_key, loss_key = jax.random.split(key, 3) # Sample from pool pool_idx, batch = pool.sample(sample_key, batch_size=batch_size, replace=False) current_state = batch["state"] current_y_idx = batch["y_idx"] current_y = y[current_y_idx] # Sort by descending loss sort_idx = jnp.argsort(jax.vmap(mse)(current_state, current_y), descending=True) pool_idx = pool_idx[sort_idx] current_state = current_state[sort_idx] current_y_idx = current_y_idx[sort_idx] # Sample a new state to replace the worst new_state, new_y_idx = sample_state(sample_state_key) current_state = current_state.at[0].set(new_state) current_y_idx = current_y_idx.at[0].set(new_y_idx) current_y = y[current_y_idx] current_z = z[current_y_idx] (loss, current_state), grad = nnx.value_and_grad( loss_fn, has_aux=True, argnums=nnx.DiffState(0, update_params) )(cs, current_state, current_z, current_y, loss_key) optimizer.update(cs, grad) pool = pool.update(pool_idx, {"state": current_state, "y_idx": current_y_idx}) return loss, pool

Main loop¶

In [15]:

num_train_steps = 8_192
print_interval = 128

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)
    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 = 8_192 print_interval = 128 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) 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")

Step    0/8192 |    3.3s | Loss 8.761e-02

Step  128/8192 |   16.7s | Loss 8.341e-02

Step  256/8192 |   30.2s | Loss 7.301e-02

Step  384/8192 |   43.5s | Loss 7.194e-02

Step  512/8192 |   57.0s | Loss 6.815e-02

Step  640/8192 |   70.4s | Loss 6.245e-02

Step  768/8192 |   83.8s | Loss 5.327e-02

Step  896/8192 |   97.2s | Loss 4.236e-02

Step 1024/8192 |  110.6s | Loss 3.558e-02

Step 1152/8192 |  124.0s | Loss 2.917e-02

Step 1280/8192 |  137.4s | Loss 2.351e-02

Step 1408/8192 |  150.8s | Loss 1.933e-02

Step 1536/8192 |  164.3s | Loss 1.565e-02

Step 1664/8192 |  177.7s | Loss 1.295e-02

Step 1792/8192 |  191.1s | Loss 1.100e-02

Step 1920/8192 |  204.5s | Loss 9.132e-03

Step 2048/8192 |  217.9s | Loss 7.979e-03

Step 2176/8192 |  231.3s | Loss 7.204e-03

Step 2304/8192 |  244.7s | Loss 6.511e-03

Step 2432/8192 |  258.1s | Loss 5.617e-03

Step 2560/8192 |  271.5s | Loss 5.272e-03

Step 2688/8192 |  284.9s | Loss 4.653e-03

Step 2816/8192 |  298.3s | Loss 4.364e-03

Step 2944/8192 |  311.6s | Loss 4.235e-03

Step 3072/8192 |  325.0s | Loss 3.963e-03

Step 3200/8192 |  338.4s | Loss 3.546e-03

Step 3328/8192 |  351.8s | Loss 3.349e-03

Step 3456/8192 |  365.2s | Loss 3.310e-03

Step 3584/8192 |  378.7s | Loss 3.063e-03

Step 3712/8192 |  392.1s | Loss 2.847e-03

Step 3840/8192 |  405.5s | Loss 2.717e-03

Step 3968/8192 |  418.9s | Loss 2.551e-03

Step 4096/8192 |  432.3s | Loss 2.440e-03

Step 4224/8192 |  445.7s | Loss 2.355e-03

Step 4352/8192 |  459.1s | Loss 2.249e-03

Step 4480/8192 |  472.5s | Loss 2.028e-03

Step 4608/8192 |  485.9s | Loss 2.050e-03

Step 4736/8192 |  499.2s | Loss 1.804e-03

Step 4864/8192 |  512.6s | Loss 1.713e-03

Step 4992/8192 |  526.0s | Loss 1.662e-03

Step 5120/8192 |  539.4s | Loss 1.623e-03

Step 5248/8192 |  552.8s | Loss 1.391e-03

Step 5376/8192 |  566.1s | Loss 1.316e-03

Step 5504/8192 |  579.5s | Loss 1.216e-03

Step 5632/8192 |  592.9s | Loss 1.125e-03

Step 5760/8192 |  606.2s | Loss 1.012e-03

Step 5888/8192 |  619.6s | Loss 9.504e-04

Step 6016/8192 |  633.0s | Loss 8.642e-04

Step 6144/8192 |  646.3s | Loss 8.209e-04

Step 6272/8192 |  659.7s | Loss 7.170e-04

Step 6400/8192 |  673.1s | Loss 6.972e-04

Step 6528/8192 |  686.4s | Loss 6.739e-04

Step 6656/8192 |  699.8s | Loss 6.256e-04

Step 6784/8192 |  713.1s | Loss 5.375e-04

Step 6912/8192 |  726.5s | Loss 5.397e-04

Step 7040/8192 |  739.8s | Loss 5.161e-04

Step 7168/8192 |  753.2s | Loss 4.671e-04

Step 7296/8192 |  766.5s | Loss 4.901e-04

Step 7424/8192 |  779.9s | Loss 4.697e-04

Step 7552/8192 |  793.2s | Loss 4.066e-04

Step 7680/8192 |  806.6s | Loss 3.711e-04

Step 7808/8192 |  819.9s | Loss 3.760e-04

Step 7936/8192 |  833.3s | Loss 3.602e-04

Step 8064/8192 |  846.6s | Loss 3.527e-04

Step 8191/8192 |  859.9s | Loss 3.625e-04
✨ Trained for 8192 steps in 860s

Run¶

In [16]:

key, subkey = jax.random.split(key)

keys = jax.random.split(subkey, y.shape[0])
state_init, _ = jax.vmap(sample_state)(keys)

state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None})
state_final, states = nnx.split_rngs(splits=y.shape[0])(
    nnx.vmap(
        lambda cs, state, z: cs(state, z, num_steps=2 * num_steps, return_states=True),
        in_axes=(state_axes, 0, 0),
    )
)(cs, state_init, z)

key, subkey = jax.random.split(key) keys = jax.random.split(subkey, y.shape[0]) state_init, _ = jax.vmap(sample_state)(keys) state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None}) state_final, states = nnx.split_rngs(splits=y.shape[0])( nnx.vmap( lambda cs, state, z: cs(state, z, num_steps=2 * num_steps, return_states=True), in_axes=(state_axes, 0, 0), ) )(cs, state_init, z)

Visualize¶

In [17]:

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

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

frames_final_rgba = nnx.vmap(
    lambda cs, state: cs.render_rgba(state),
    in_axes=(None, 0),
)(cs, state_final)

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

In [18]:

mediapy.show_images(frames_final_rgba.repeat(2, axis=-3).repeat(2, axis=-2))
mediapy.show_videos(frames.repeat(2, axis=-3).repeat(2, axis=-2))

mediapy.show_images(frames_final_rgba.repeat(2, axis=-3).repeat(2, axis=-2)) mediapy.show_videos(frames.repeat(2, axis=-3).repeat(2, axis=-2))

Interpolation¶

Dog-Panda¶

In [19]:

# Define latent encoding
z_dog_panda = jnp.array([0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.5])

# Sample initial state
key, subkey = jax.random.split(key)
state_init, _ = sample_state(subkey)

# Run
state_final, states = cs(
    state_init, z_dog_panda, num_steps=2 * num_steps, return_states=True
)


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

mediapy.show_video(frames.repeat(2, axis=-3).repeat(2, axis=-2))

# Define latent encoding z_dog_panda = jnp.array([0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.5]) # Sample initial state key, subkey = jax.random.split(key) state_init, _ = sample_state(subkey) # Run state_final, states = cs( state_init, z_dog_panda, num_steps=2 * num_steps, return_states=True ) # Visualize states = jnp.concatenate([state_init[None], states]) frames = nnx.vmap( lambda cs, state: cs.render(state), in_axes=(None, 0), )(cs, states) mediapy.show_video(frames.repeat(2, axis=-3).repeat(2, axis=-2))

Fox-Panda¶

In [20]:

# Define latent encoding
z_fox_panda = jnp.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.5, 0.0, 0.5])

# Sample initial state
key, subkey = jax.random.split(key)
state_init, _ = sample_state(subkey)

# Run
state_final, states = cs(
    state_init, z_fox_panda, num_steps=2 * num_steps, return_states=True
)


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

mediapy.show_video(frames.repeat(2, axis=-3).repeat(2, axis=-2))

# Define latent encoding z_fox_panda = jnp.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.5, 0.0, 0.5]) # Sample initial state key, subkey = jax.random.split(key) state_init, _ = sample_state(subkey) # Run state_final, states = cs( state_init, z_fox_panda, num_steps=2 * num_steps, return_states=True ) # Visualize states = jnp.concatenate([state_init[None], states]) frames = nnx.vmap( lambda cs, state: cs.render(state), in_axes=(None, 0), )(cs, states) mediapy.show_video(frames.repeat(2, axis=-3).repeat(2, axis=-2))

Cat-Panda¶

In [21]:

# Define two target images
cat_idx = 1
panda_idx = 7

# Interpolate between the two latent encodings
num_interpolations = 8
alphas = jnp.linspace(0.0, 1.0, num_interpolations)
z_interpolation = jnp.array(
    [(1.0 - alpha) * z[cat_idx] + alpha * z[panda_idx] for alpha in alphas]
)

# Sample initial state
key, subkey = jax.random.split(key)
state_init, _ = sample_state(subkey)

# Run
state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None})
final_states = nnx.split_rngs(splits=alphas.shape[0])(
    nnx.vmap(
        lambda cs, state, z: cs(state, z, num_steps=2 * num_steps),
        in_axes=(state_axes, None, 0),
    )
)(cs, state_init, z_interpolation)

# Define two target images cat_idx = 1 panda_idx = 7 # Interpolate between the two latent encodings num_interpolations = 8 alphas = jnp.linspace(0.0, 1.0, num_interpolations) z_interpolation = jnp.array( [(1.0 - alpha) * z[cat_idx] + alpha * z[panda_idx] for alpha in alphas] ) # Sample initial state key, subkey = jax.random.split(key) state_init, _ = sample_state(subkey) # Run state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None}) final_states = nnx.split_rngs(splits=alphas.shape[0])( nnx.vmap( lambda cs, state, z: cs(state, z, num_steps=2 * num_steps), in_axes=(state_axes, None, 0), ) )(cs, state_init, z_interpolation)

In [22]:

frames = nnx.vmap(
    lambda cs, state: cs.render_rgba(state),
    in_axes=(None, 0),
)(cs, final_states)

mediapy.show_images(frames.repeat(2, axis=-3).repeat(2, axis=-2))

frames = nnx.vmap( lambda cs, state: cs.render_rgba(state), in_axes=(None, 0), )(cs, final_states) mediapy.show_images(frames.repeat(2, axis=-3).repeat(2, axis=-2))