Skip to content

Update

cax.core.update.update

Update base module.

Update

Bases: Module

Base class for update modules.

Subclasses implement transforms mapping a state and a perception (and optional input) to the next state.

Source code in src/cax/core/update/update.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class Update[State, Perception, Input](nnx.Module):
    """Base class for update modules.

    Subclasses implement transforms mapping a state and a perception (and optional
    input) to the next state.
    """

    def __call__(
        self, state: State, perception: Perception, input: Input | None = None
    ) -> State:
        """Process the current state, perception, and input to produce a new state.

        This method should be implemented by subclasses to define specific update logic.

        Args:
            state: Current state.
            perception: Current perception.
            input: Optional input.

        Returns:
            Next state.

        """
        raise NotImplementedError

__call__(state, perception, input=None)

Process the current state, perception, and input to produce a new state.

This method should be implemented by subclasses to define specific update logic.

Parameters:

Name Type Description Default
state State

Current state.

 required
perception Perception

Current perception.

 required
input Input | None

Optional input.

 None

Returns:

Type Description
State

Next state.
Source code in src/cax/core/update/update.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
def __call__(
    self, state: State, perception: Perception, input: Input | None = None
) -> State:
    """Process the current state, perception, and input to produce a new state.

    This method should be implemented by subclasses to define specific update logic.

    Args:
        state: Current state.
        perception: Current perception.
        input: Optional input.

    Returns:
        Next state.

    """
    raise NotImplementedError

cax.core.update.mlp_update

MLP update module.

MLPUpdate

Bases: Update[Array, Array, Array]

MLP update class.

Maps a perception (and optional input) to the next state using pointwise convolutional layers (kernel size 1) applied independently at each spatial position.

Source code in src/cax/core/update/mlp_update.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
class MLPUpdate(Update[Array, Array, Array]):
    """MLP update class.

    Maps a perception (and optional input) to the next state using pointwise
    convolutional layers (kernel size 1) applied independently at each spatial position.

    """

    def __init__(
        self,
        *,
        num_spatial_dims: int,
        channel_size: int,
        perception_size: int,
        hidden_layer_sizes: tuple[int, ...],
        activation_fn: Callable[[Array], Array] = nnx.relu,
        zeros_init: bool = False,
        rngs: nnx.Rngs,
    ):
        """Initialize MLP update.

        Args:
            num_spatial_dims: Number of spatial dimensions.
            channel_size: Number of channels in the output.
            perception_size: Size of the input perception.
            hidden_layer_sizes: Sizes of hidden layers.
            activation_fn: Activation function to use.
            zeros_init: Whether to use zeros initialization for the weights of the last
                layer.
            rngs: rng key.

        """
        in_features = (perception_size, *hidden_layer_sizes)
        out_features = (*hidden_layer_sizes, channel_size)
        kernel_init = [default_kernel_init for _ in hidden_layer_sizes] + [
            initializers.zeros_init() if zeros_init else default_kernel_init
        ]
        self.layers = nnx.List(
            [
                nnx.Conv(
                    in_features,
                    out_features,
                    kernel_size=num_spatial_dims * (1,),
                    kernel_init=kernel_init,
                    rngs=rngs,
                )
                for in_features, out_features, kernel_init in zip(
                    in_features, out_features, kernel_init, strict=True
                )
            ]
        )
        self.activation_fn = activation_fn

    @override
    def __call__(
        self, state: Array, perception: Array, input: Array | None = None
    ) -> Array:
        """Process the current state, perception, and input to produce a new state.

        If input is provided, it is concatenated to the perception along the channel
        axis before being passed through the layers.

        Args:
            state: Current state.
            perception: Current perception.
            input: Optional input.

        Returns:
            Next state.

        """
        if input is not None:
            perception = jnp.concatenate([perception, input], axis=-1)

        for layer in self.layers[:-1]:
            perception = self.activation_fn(layer(perception))
        state = self.layers[-1](perception)
        return state

__init__(*, num_spatial_dims, channel_size, perception_size, hidden_layer_sizes, activation_fn=nnx.relu, zeros_init=False, rngs)

Initialize MLP update.

Parameters:

Name Type Description Default
num_spatial_dims int

Number of spatial dimensions.

 required
channel_size int

Number of channels in the output.

 required
perception_size int

Size of the input perception.

 required
hidden_layer_sizes tuple[int, ...]

Sizes of hidden layers.

 required
activation_fn Callable[[Array], Array]

Activation function to use.

 relu
zeros_init bool

Whether to use zeros initialization for the weights of the last layer.

 False
rngs Rngs

rng key.

 required
Source code in src/cax/core/update/mlp_update.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def __init__(
    self,
    *,
    num_spatial_dims: int,
    channel_size: int,
    perception_size: int,
    hidden_layer_sizes: tuple[int, ...],
    activation_fn: Callable[[Array], Array] = nnx.relu,
    zeros_init: bool = False,
    rngs: nnx.Rngs,
):
    """Initialize MLP update.

    Args:
        num_spatial_dims: Number of spatial dimensions.
        channel_size: Number of channels in the output.
        perception_size: Size of the input perception.
        hidden_layer_sizes: Sizes of hidden layers.
        activation_fn: Activation function to use.
        zeros_init: Whether to use zeros initialization for the weights of the last
            layer.
        rngs: rng key.

    """
    in_features = (perception_size, *hidden_layer_sizes)
    out_features = (*hidden_layer_sizes, channel_size)
    kernel_init = [default_kernel_init for _ in hidden_layer_sizes] + [
        initializers.zeros_init() if zeros_init else default_kernel_init
    ]
    self.layers = nnx.List(
        [
            nnx.Conv(
                in_features,
                out_features,
                kernel_size=num_spatial_dims * (1,),
                kernel_init=kernel_init,
                rngs=rngs,
            )
            for in_features, out_features, kernel_init in zip(
                in_features, out_features, kernel_init, strict=True
            )
        ]
    )
    self.activation_fn = activation_fn

__call__(state, perception, input=None)

Process the current state, perception, and input to produce a new state.

If input is provided, it is concatenated to the perception along the channel axis before being passed through the layers.

Parameters:

Name Type Description Default
state Array

Current state.

 required
perception Array

Current perception.

 required
input Array | None

Optional input.

 None

Returns:

Type Description
Array

Next state.
Source code in src/cax/core/update/mlp_update.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
@override
def __call__(
    self, state: Array, perception: Array, input: Array | None = None
) -> Array:
    """Process the current state, perception, and input to produce a new state.

    If input is provided, it is concatenated to the perception along the channel
    axis before being passed through the layers.

    Args:
        state: Current state.
        perception: Current perception.
        input: Optional input.

    Returns:
        Next state.

    """
    if input is not None:
        perception = jnp.concatenate([perception, input], axis=-1)

    for layer in self.layers[:-1]:
        perception = self.activation_fn(layer(perception))
    state = self.layers[-1](perception)
    return state

cax.core.update.residual_update

Residual update module.

ResidualUpdate

Bases: MLPUpdate

Residual update class.

Extends the MLP update with a residual connection and cell dropout applied to the update.

Source code in src/cax/core/update/residual_update.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
class ResidualUpdate(MLPUpdate):
    """Residual update class.

    Extends the MLP update with a residual connection and cell dropout applied to the
    update.
    """

    def __init__(
        self,
        *,
        num_spatial_dims: int,
        channel_size: int,
        perception_size: int,
        hidden_layer_sizes: tuple[int, ...],
        activation_fn: Callable[[Array], Array] = nnx.relu,
        step_size: float = 1.0,
        cell_dropout_rate: float = 0.0,
        zeros_init: bool = False,
        rngs: nnx.Rngs,
    ):
        """Initialize the ResidualUpdate module.

        Args:
            num_spatial_dims: Number of spatial dimensions.
            channel_size: Number of channels in the state.
            perception_size: Size of the perception input.
            hidden_layer_sizes: Sizes of hidden layers in the MLP.
            activation_fn: Activation function to use.
            step_size: Step size for the residual update.
            cell_dropout_rate: Dropout rate for cell updates.
            zeros_init: Whether to use zeros initialization for the weights of the last
                layer.
            rngs: rng key.

        """
        super().__init__(
            num_spatial_dims=num_spatial_dims,
            channel_size=channel_size,
            perception_size=perception_size,
            hidden_layer_sizes=hidden_layer_sizes,
            activation_fn=activation_fn,
            zeros_init=zeros_init,
            rngs=rngs,
        )
        self.dropout = nnx.Dropout(
            rate=cell_dropout_rate, broadcast_dims=(-1,), rngs=rngs
        )
        self.step_size = step_size

    @override
    def __call__(
        self, state: Array, perception: Array, input: Array | None = None
    ) -> Array:
        """Process the current state, perception, and input to produce a new state.

        Args:
            state: Current state.
            perception: Current perception.
            input: Optional input.

        Returns:
            Next state.

        """
        update = super().__call__(state, perception, input)
        update = self.dropout(update)
        state += self.step_size * update
        return state

__init__(*, num_spatial_dims, channel_size, perception_size, hidden_layer_sizes, activation_fn=nnx.relu, step_size=1.0, cell_dropout_rate=0.0, zeros_init=False, rngs)

Initialize the ResidualUpdate module.

Parameters:

Name Type Description Default
num_spatial_dims int

Number of spatial dimensions.

 required
channel_size int

Number of channels in the state.

 required
perception_size int

Size of the perception input.

 required
hidden_layer_sizes tuple[int, ...]

Sizes of hidden layers in the MLP.

 required
activation_fn Callable[[Array], Array]

Activation function to use.

 relu
step_size float

Step size for the residual update.

 1.0
cell_dropout_rate float

Dropout rate for cell updates.

 0.0
zeros_init bool

Whether to use zeros initialization for the weights of the last layer.

 False
rngs Rngs

rng key.

 required
Source code in src/cax/core/update/residual_update.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def __init__(
    self,
    *,
    num_spatial_dims: int,
    channel_size: int,
    perception_size: int,
    hidden_layer_sizes: tuple[int, ...],
    activation_fn: Callable[[Array], Array] = nnx.relu,
    step_size: float = 1.0,
    cell_dropout_rate: float = 0.0,
    zeros_init: bool = False,
    rngs: nnx.Rngs,
):
    """Initialize the ResidualUpdate module.

    Args:
        num_spatial_dims: Number of spatial dimensions.
        channel_size: Number of channels in the state.
        perception_size: Size of the perception input.
        hidden_layer_sizes: Sizes of hidden layers in the MLP.
        activation_fn: Activation function to use.
        step_size: Step size for the residual update.
        cell_dropout_rate: Dropout rate for cell updates.
        zeros_init: Whether to use zeros initialization for the weights of the last
            layer.
        rngs: rng key.

    """
    super().__init__(
        num_spatial_dims=num_spatial_dims,
        channel_size=channel_size,
        perception_size=perception_size,
        hidden_layer_sizes=hidden_layer_sizes,
        activation_fn=activation_fn,
        zeros_init=zeros_init,
        rngs=rngs,
    )
    self.dropout = nnx.Dropout(
        rate=cell_dropout_rate, broadcast_dims=(-1,), rngs=rngs
    )
    self.step_size = step_size

__call__(state, perception, input=None)

Process the current state, perception, and input to produce a new state.

Parameters:

Name Type Description Default
state Array

Current state.

 required
perception Array

Current perception.

 required
input Array | None

Optional input.

 None

Returns:

Type Description
Array

Next state.
Source code in src/cax/core/update/residual_update.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
@override
def __call__(
    self, state: Array, perception: Array, input: Array | None = None
) -> Array:
    """Process the current state, perception, and input to produce a new state.

    Args:
        state: Current state.
        perception: Current perception.
        input: Optional input.

    Returns:
        Next state.

    """
    update = super().__call__(state, perception, input)
    update = self.dropout(update)
    state += self.step_size * update
    return state

cax.core.update.nca_update

Neural Cellular Automata update module.

NCAUpdate

Bases: ResidualUpdate

Neural Cellular Automata update class.

Builds on the residual update and applies an alive mask so that only active cells update.

Source code in src/cax/core/update/nca_update.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
class NCAUpdate(ResidualUpdate):
    """Neural Cellular Automata update class.

    Builds on the residual update and applies an alive mask so that only active cells
    update.
    """

    def __init__(
        self,
        *,
        channel_size: int,
        perception_size: int,
        hidden_layer_sizes: tuple[int, ...],
        activation_fn: Callable[[Array], Array] = nnx.relu,
        step_size: float = 1.0,
        cell_dropout_rate: float = 0.0,
        kernel_size: Sequence[int] = (3, 3),
        alive_threshold: float = 0.1,
        zeros_init: bool = False,
        rngs: nnx.Rngs,
    ):
        """Initialize NCA update.

        Args:
            channel_size: Number of input channels.
            perception_size: Size of the perception.
            hidden_layer_sizes: Sizes of hidden layers.
            activation_fn: Activation function to use.
            step_size: Step size for the update.
            cell_dropout_rate: Dropout rate for cells.
            kernel_size: Size of the convolutional kernel.
            alive_threshold: Threshold for determining if a cell is alive.
            zeros_init: Whether to use zeros initialization for the weights of the last
                layer.
            rngs: rng key.

        """
        super().__init__(
            num_spatial_dims=len(kernel_size),
            channel_size=channel_size,
            perception_size=perception_size,
            hidden_layer_sizes=hidden_layer_sizes,
            rngs=rngs,
            activation_fn=activation_fn,
            step_size=step_size,
            cell_dropout_rate=cell_dropout_rate,
            zeros_init=zeros_init,
        )
        self.pool = partial(nnx.max_pool, window_shape=kernel_size, padding="SAME")
        self.alive_threshold = alive_threshold

    @override
    def __call__(
        self, state: Array, perception: Array, input: Array | None = None
    ) -> Array:
        """Process the current state, perception, and input to produce a new state.

        Args:
            state: Current state.
            perception: Current perception.
            input: Optional input.

        Returns:
            Next state.

        """
        alive_mask = self.get_alive_mask(state)
        state = super().__call__(state, perception, input)
        alive_mask &= self.get_alive_mask(state)
        return alive_mask * state

    def get_alive_mask(self, state: Array) -> Array:
        """Generate a mask of alive cells based on the current state.

        Args:
            state: Current state.

        Returns:
            A boolean mask indicating which cells are alive.

        """
        alive = state_to_alive(state)
        alive_mask: Array = self.pool(alive) > self.alive_threshold
        return alive_mask

__init__(*, channel_size, perception_size, hidden_layer_sizes, activation_fn=nnx.relu, step_size=1.0, cell_dropout_rate=0.0, kernel_size=(3, 3), alive_threshold=0.1, zeros_init=False, rngs)

Initialize NCA update.

Parameters:

Name Type Description Default
channel_size int

Number of input channels.

 required
perception_size int

Size of the perception.

 required
hidden_layer_sizes tuple[int, ...]

Sizes of hidden layers.

 required
activation_fn Callable[[Array], Array]

Activation function to use.

 relu
step_size float

Step size for the update.

 1.0
cell_dropout_rate float

Dropout rate for cells.

 0.0
kernel_size Sequence[int]

Size of the convolutional kernel.

 (3, 3)
alive_threshold float

Threshold for determining if a cell is alive.

 0.1
zeros_init bool

Whether to use zeros initialization for the weights of the last layer.

 False
rngs Rngs

rng key.

 required
Source code in src/cax/core/update/nca_update.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def __init__(
    self,
    *,
    channel_size: int,
    perception_size: int,
    hidden_layer_sizes: tuple[int, ...],
    activation_fn: Callable[[Array], Array] = nnx.relu,
    step_size: float = 1.0,
    cell_dropout_rate: float = 0.0,
    kernel_size: Sequence[int] = (3, 3),
    alive_threshold: float = 0.1,
    zeros_init: bool = False,
    rngs: nnx.Rngs,
):
    """Initialize NCA update.

    Args:
        channel_size: Number of input channels.
        perception_size: Size of the perception.
        hidden_layer_sizes: Sizes of hidden layers.
        activation_fn: Activation function to use.
        step_size: Step size for the update.
        cell_dropout_rate: Dropout rate for cells.
        kernel_size: Size of the convolutional kernel.
        alive_threshold: Threshold for determining if a cell is alive.
        zeros_init: Whether to use zeros initialization for the weights of the last
            layer.
        rngs: rng key.

    """
    super().__init__(
        num_spatial_dims=len(kernel_size),
        channel_size=channel_size,
        perception_size=perception_size,
        hidden_layer_sizes=hidden_layer_sizes,
        rngs=rngs,
        activation_fn=activation_fn,
        step_size=step_size,
        cell_dropout_rate=cell_dropout_rate,
        zeros_init=zeros_init,
    )
    self.pool = partial(nnx.max_pool, window_shape=kernel_size, padding="SAME")
    self.alive_threshold = alive_threshold

__call__(state, perception, input=None)

Process the current state, perception, and input to produce a new state.

Parameters:

Name Type Description Default
state Array

Current state.

 required
perception Array

Current perception.

 required
input Array | None

Optional input.

 None

Returns:

Type Description
Array

Next state.
Source code in src/cax/core/update/nca_update.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
@override
def __call__(
    self, state: Array, perception: Array, input: Array | None = None
) -> Array:
    """Process the current state, perception, and input to produce a new state.

    Args:
        state: Current state.
        perception: Current perception.
        input: Optional input.

    Returns:
        Next state.

    """
    alive_mask = self.get_alive_mask(state)
    state = super().__call__(state, perception, input)
    alive_mask &= self.get_alive_mask(state)
    return alive_mask * state

get_alive_mask(state)

Generate a mask of alive cells based on the current state.

Parameters:

Name Type Description Default
state Array

Current state.

 required

Returns:

Type Description
Array

A boolean mask indicating which cells are alive.
Source code in src/cax/core/update/nca_update.py 
84
85
86
87
88
89
90
91
92
93
94
95
96
def get_alive_mask(self, state: Array) -> Array:
    """Generate a mask of alive cells based on the current state.

    Args:
        state: Current state.

    Returns:
        A boolean mask indicating which cells are alive.

    """
    alive = state_to_alive(state)
    alive_mask: Array = self.pool(alive) > self.alive_threshold
    return alive_mask

state_to_alive(state)

Extract the 'alive' component from the state.

Parameters:

Name Type Description Default
state Array

Input state.

 required

Returns:

Type Description
Array

The 'alive' component of the state.
Source code in src/cax/core/update/nca_update.py 
 99
100
101
102
103
104
105
106
107
108
109
def state_to_alive(state: Array) -> Array:
    """Extract the 'alive' component from the state.

    Args:
        state: Input state.

    Returns:
        The 'alive' component of the state.

    """
    return state[..., -1:]