Skip to content

Perceive

cax.core.perceive.perceive

Perceive base module.

Perceive

Bases: Module

Base class for perception modules.

Subclasses implement neighborhood gathering or convolutional transforms that map a state to a perception. Perceptions are PyTrees — commonly arrays shaped (..., *spatial_dims, perception_size), or a dataclass of arrays — and the type parameter names which one a subclass produces, so the perceive-update handoff is checked instead of erased.

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

    Subclasses implement neighborhood gathering or convolutional transforms that map a
    state to a perception. Perceptions are PyTrees — commonly arrays shaped
    `(..., *spatial_dims, perception_size)`, or a dataclass of arrays — and the type
    parameter names which one a subclass produces, so the perceive-update handoff is
    checked instead of erased.
    """

    def __call__(self, state: State) -> Perception:
        """Process the current state to produce a perception.

        Args:
            state: Current state.

        Returns:
            Perception derived from `state`.

        """
        raise NotImplementedError

__call__(state)

Process the current state to produce a perception.

Parameters:

Name Type Description Default
state State

Current state.

 required

Returns:

Type Description
Perception

Perception derived from state.
Source code in src/cax/core/perceive/perceive.py 
16
17
18
19
20
21
22
23
24
25
26
def __call__(self, state: State) -> Perception:
    """Process the current state to produce a perception.

    Args:
        state: Current state.

    Returns:
        Perception derived from `state`.

    """
    raise NotImplementedError

cax.core.perceive.neighborhood_perceive

Neighborhood perceive base module.

This module provides the shared pad-and-slice implementation for neighborhood-based perception. Both Moore and Von Neumann perceive modules inherit from this base class, differing only in which neighbor shifts they enumerate.

NeighborhoodPerceive

Bases: Perceive[Array, Array]

Base class for neighborhood-based perception using pad-and-slice.

Subclasses implement _get_shifts to define which neighbor offsets belong to the neighborhood. This class handles padding, slicing, and optional reduction over neighbors.

Source code in src/cax/core/perceive/neighborhood_perceive.py 
 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
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
class NeighborhoodPerceive(Perceive[Array, Array]):
    """Base class for neighborhood-based perception using pad-and-slice.

    Subclasses implement `_get_shifts` to define which neighbor offsets belong to the
    neighborhood. This class handles padding, slicing, and optional reduction over
    neighbors.
    """

    def __init__(
        self,
        *,
        num_spatial_dims: int,
        radius: int,
        padding: str = "CIRCULAR",
        include_center: bool = True,
        reduce_fn: Callable[..., Array] | None = None,
    ):
        """Initialize neighborhood perceive.

        Args:
            num_spatial_dims: Number of spatial dimensions.
            radius: Neighborhood radius.
            padding: Boundary condition mode. One of "CIRCULAR" (periodic/torus),
                "ZERO" (zero-padded), "REFLECT" (mirror), or "EDGE" (clamp to boundary).
            include_center: Whether to include the center cell in the output.
            reduce_fn: Optional reduction function applied over the neighbor axis. If
                None, neighbors are concatenated along the channel axis. If provided, it
                is called as `reduce_fn(stacked_neighbors, axis=0)` and the result is
                concatenated with the center (if `include_center` is True).

        """
        if padding not in _PADDING_TO_PAD_MODE:
            valid = ", ".join(sorted(_PADDING_TO_PAD_MODE.keys()))
            raise ValueError(f"padding must be one of {valid}, got {padding!r}")

        self.num_spatial_dims = num_spatial_dims
        self.radius = radius
        self.padding = padding
        self.include_center = include_center
        self.reduce_fn = reduce_fn

    def _get_shifts(self) -> list[tuple[int, ...]]:
        """Return the list of neighbor shifts (excluding center).

        Each shift is a tuple of length `num_spatial_dims` with integer offsets
        in the range `[-radius, radius]`.

        Returns:
            List of shift tuples defining the neighborhood.

        """
        raise NotImplementedError

    @override
    def __call__(self, state: Array) -> Array:
        """Apply neighborhood perception to the input state.

        The input is assumed to have shape `(..., *spatial_dims, channel_size)`.

        When `reduce_fn` is None, neighbors are concatenated along the channel axis:
            - With `include_center=True`: shape
              `(..., *spatial, channel_size * (N + 1))`
            - With `include_center=False`: shape `(..., *spatial, channel_size * N)`
            where N is the number of neighbor shifts.

        When `reduce_fn` is provided, neighbors are stacked and reduced:
            - With `include_center=True`: shape `(..., *spatial, 2 * channel_size)`
            - With `include_center=False`: shape `(..., *spatial, channel_size)`

        Args:
            state: State of the cellular automaton.

        Returns:
            Neighborhood perception.

        """
        spatial_start = state.ndim - self.num_spatial_dims - 1
        spatial_sizes = state.shape[
            spatial_start : spatial_start + self.num_spatial_dims
        ]

        pad_mode = _PADDING_TO_PAD_MODE[self.padding]
        pad_widths = (
            [(0, 0)] * spatial_start
            + [(self.radius, self.radius)] * self.num_spatial_dims
            + [(0, 0)]
        )
        padded = jnp.pad(state, pad_widths, mode=pad_mode)

        shifts = self._get_shifts()

        neighbors = []
        for shift in shifts:
            slices = [slice(None)] * spatial_start
            for dim_idx in range(self.num_spatial_dims):
                start = self.radius - shift[dim_idx]
                slices.append(slice(start, start + spatial_sizes[dim_idx]))
            slices.append(slice(None))
            neighbors.append(padded[tuple(slices)])

        if self.reduce_fn is not None:
            stacked = jnp.stack(neighbors, axis=0)
            reduced = self.reduce_fn(stacked, axis=0)
            if self.include_center:
                return jnp.concatenate([state, reduced], axis=-1)
            return reduced

        if self.include_center:
            return jnp.concatenate([state, *neighbors], axis=-1)
        return jnp.concatenate(neighbors, axis=-1)

__init__(*, num_spatial_dims, radius, padding='CIRCULAR', include_center=True, reduce_fn=None)

Initialize neighborhood perceive.

Parameters:

Name Type Description Default
num_spatial_dims int

Number of spatial dimensions.

 required
radius int

Neighborhood radius.

 required
padding str

Boundary condition mode. One of "CIRCULAR" (periodic/torus), "ZERO" (zero-padded), "REFLECT" (mirror), or "EDGE" (clamp to boundary).

 'CIRCULAR'
include_center bool

Whether to include the center cell in the output.

 True
reduce_fn Callable[..., Array] | None

Optional reduction function applied over the neighbor axis. If None, neighbors are concatenated along the channel axis. If provided, it is called as reduce_fn(stacked_neighbors, axis=0) and the result is concatenated with the center (if include_center is True).

 None
Source code in src/cax/core/perceive/neighborhood_perceive.py 
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
def __init__(
    self,
    *,
    num_spatial_dims: int,
    radius: int,
    padding: str = "CIRCULAR",
    include_center: bool = True,
    reduce_fn: Callable[..., Array] | None = None,
):
    """Initialize neighborhood perceive.

    Args:
        num_spatial_dims: Number of spatial dimensions.
        radius: Neighborhood radius.
        padding: Boundary condition mode. One of "CIRCULAR" (periodic/torus),
            "ZERO" (zero-padded), "REFLECT" (mirror), or "EDGE" (clamp to boundary).
        include_center: Whether to include the center cell in the output.
        reduce_fn: Optional reduction function applied over the neighbor axis. If
            None, neighbors are concatenated along the channel axis. If provided, it
            is called as `reduce_fn(stacked_neighbors, axis=0)` and the result is
            concatenated with the center (if `include_center` is True).

    """
    if padding not in _PADDING_TO_PAD_MODE:
        valid = ", ".join(sorted(_PADDING_TO_PAD_MODE.keys()))
        raise ValueError(f"padding must be one of {valid}, got {padding!r}")

    self.num_spatial_dims = num_spatial_dims
    self.radius = radius
    self.padding = padding
    self.include_center = include_center
    self.reduce_fn = reduce_fn

__call__(state)

Apply neighborhood perception to the input state.

The input is assumed to have shape (..., *spatial_dims, channel_size).

When reduce_fn is None, neighbors are concatenated along the channel axis: - With include_center=True: shape (..., *spatial, channel_size * (N + 1)) - With include_center=False: shape (..., *spatial, channel_size * N) where N is the number of neighbor shifts.

When reduce_fn is provided, neighbors are stacked and reduced: - With include_center=True: shape (..., *spatial, 2 * channel_size) - With include_center=False: shape (..., *spatial, channel_size)

Parameters:

Name Type Description Default
state Array

State of the cellular automaton.

 required

Returns:

Type Description
Array

Neighborhood perception.
Source code in src/cax/core/perceive/neighborhood_perceive.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
@override
def __call__(self, state: Array) -> Array:
    """Apply neighborhood perception to the input state.

    The input is assumed to have shape `(..., *spatial_dims, channel_size)`.

    When `reduce_fn` is None, neighbors are concatenated along the channel axis:
        - With `include_center=True`: shape
          `(..., *spatial, channel_size * (N + 1))`
        - With `include_center=False`: shape `(..., *spatial, channel_size * N)`
        where N is the number of neighbor shifts.

    When `reduce_fn` is provided, neighbors are stacked and reduced:
        - With `include_center=True`: shape `(..., *spatial, 2 * channel_size)`
        - With `include_center=False`: shape `(..., *spatial, channel_size)`

    Args:
        state: State of the cellular automaton.

    Returns:
        Neighborhood perception.

    """
    spatial_start = state.ndim - self.num_spatial_dims - 1
    spatial_sizes = state.shape[
        spatial_start : spatial_start + self.num_spatial_dims
    ]

    pad_mode = _PADDING_TO_PAD_MODE[self.padding]
    pad_widths = (
        [(0, 0)] * spatial_start
        + [(self.radius, self.radius)] * self.num_spatial_dims
        + [(0, 0)]
    )
    padded = jnp.pad(state, pad_widths, mode=pad_mode)

    shifts = self._get_shifts()

    neighbors = []
    for shift in shifts:
        slices = [slice(None)] * spatial_start
        for dim_idx in range(self.num_spatial_dims):
            start = self.radius - shift[dim_idx]
            slices.append(slice(start, start + spatial_sizes[dim_idx]))
        slices.append(slice(None))
        neighbors.append(padded[tuple(slices)])

    if self.reduce_fn is not None:
        stacked = jnp.stack(neighbors, axis=0)
        reduced = self.reduce_fn(stacked, axis=0)
        if self.include_center:
            return jnp.concatenate([state, reduced], axis=-1)
        return reduced

    if self.include_center:
        return jnp.concatenate([state, *neighbors], axis=-1)
    return jnp.concatenate(neighbors, axis=-1)

cax.core.perceive.moore_perceive

Moore perceive module.

MoorePerceive

Bases: NeighborhoodPerceive

Moore perceive class.

This class implements perception based on the Moore neighborhood. The Moore neighborhood includes all cells within Chebyshev distance radius of the central cell — i.e., the full hypercube of side length 2 * radius + 1 excluding the center.

Source code in src/cax/core/perceive/moore_perceive.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
class MoorePerceive(NeighborhoodPerceive):
    """Moore perceive class.

    This class implements perception based on the Moore neighborhood.
    The Moore neighborhood includes all cells within Chebyshev distance `radius` of the
    central cell — i.e., the full hypercube of side length `2 * radius + 1` excluding
    the center.
    """

    def __init__(
        self,
        *,
        num_spatial_dims: int,
        radius: int,
        padding: str = "CIRCULAR",
        include_center: bool = True,
        reduce_fn: Callable[..., Array] | None = None,
    ):
        """Initialize Moore perceive.

        Args:
            num_spatial_dims: Number of spatial dimensions.
            radius: Chebyshev distance defining the Moore neighborhood extent.
            padding: Boundary condition mode. One of "CIRCULAR" (periodic/torus),
                "ZERO" (zero-padded), "REFLECT" (mirror), or "EDGE" (clamp to boundary).
            include_center: Whether to include the center cell in the output.
            reduce_fn: Optional reduction function applied over the neighbor axis. If
                None, neighbors are concatenated along the channel axis. If provided, it
                is called as `reduce_fn(stacked_neighbors, axis=0)` and the result is
                concatenated with the center (if `include_center` is True).

        """
        super().__init__(
            num_spatial_dims=num_spatial_dims,
            radius=radius,
            padding=padding,
            include_center=include_center,
            reduce_fn=reduce_fn,
        )

    @override
    def _get_shifts(self) -> list[tuple[int, ...]]:
        """Return all shifts in the Moore neighborhood (excluding center).

        Returns:
            List of shift tuples covering the full hypercube minus the origin.

        """
        return [
            shift
            for shift in product(
                range(-self.radius, self.radius + 1), repeat=self.num_spatial_dims
            )
            if shift != (0,) * self.num_spatial_dims
        ]

__init__(*, num_spatial_dims, radius, padding='CIRCULAR', include_center=True, reduce_fn=None)

Initialize Moore perceive.

Parameters:

Name Type Description Default
num_spatial_dims int

Number of spatial dimensions.

 required
radius int

Chebyshev distance defining the Moore neighborhood extent.

 required
padding str

Boundary condition mode. One of "CIRCULAR" (periodic/torus), "ZERO" (zero-padded), "REFLECT" (mirror), or "EDGE" (clamp to boundary).

 'CIRCULAR'
include_center bool

Whether to include the center cell in the output.

 True
reduce_fn Callable[..., Array] | None

Optional reduction function applied over the neighbor axis. If None, neighbors are concatenated along the channel axis. If provided, it is called as reduce_fn(stacked_neighbors, axis=0) and the result is concatenated with the center (if include_center is True).

 None
Source code in src/cax/core/perceive/moore_perceive.py 
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
def __init__(
    self,
    *,
    num_spatial_dims: int,
    radius: int,
    padding: str = "CIRCULAR",
    include_center: bool = True,
    reduce_fn: Callable[..., Array] | None = None,
):
    """Initialize Moore perceive.

    Args:
        num_spatial_dims: Number of spatial dimensions.
        radius: Chebyshev distance defining the Moore neighborhood extent.
        padding: Boundary condition mode. One of "CIRCULAR" (periodic/torus),
            "ZERO" (zero-padded), "REFLECT" (mirror), or "EDGE" (clamp to boundary).
        include_center: Whether to include the center cell in the output.
        reduce_fn: Optional reduction function applied over the neighbor axis. If
            None, neighbors are concatenated along the channel axis. If provided, it
            is called as `reduce_fn(stacked_neighbors, axis=0)` and the result is
            concatenated with the center (if `include_center` is True).

    """
    super().__init__(
        num_spatial_dims=num_spatial_dims,
        radius=radius,
        padding=padding,
        include_center=include_center,
        reduce_fn=reduce_fn,
    )

__call__(state)

Apply neighborhood perception to the input state.

The input is assumed to have shape (..., *spatial_dims, channel_size).

When reduce_fn is None, neighbors are concatenated along the channel axis: - With include_center=True: shape (..., *spatial, channel_size * (N + 1)) - With include_center=False: shape (..., *spatial, channel_size * N) where N is the number of neighbor shifts.

When reduce_fn is provided, neighbors are stacked and reduced: - With include_center=True: shape (..., *spatial, 2 * channel_size) - With include_center=False: shape (..., *spatial, channel_size)

Parameters:

Name Type Description Default
state Array

State of the cellular automaton.

 required

Returns:

Type Description
Array

Neighborhood perception.
Source code in src/cax/core/perceive/neighborhood_perceive.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
@override
def __call__(self, state: Array) -> Array:
    """Apply neighborhood perception to the input state.

    The input is assumed to have shape `(..., *spatial_dims, channel_size)`.

    When `reduce_fn` is None, neighbors are concatenated along the channel axis:
        - With `include_center=True`: shape
          `(..., *spatial, channel_size * (N + 1))`
        - With `include_center=False`: shape `(..., *spatial, channel_size * N)`
        where N is the number of neighbor shifts.

    When `reduce_fn` is provided, neighbors are stacked and reduced:
        - With `include_center=True`: shape `(..., *spatial, 2 * channel_size)`
        - With `include_center=False`: shape `(..., *spatial, channel_size)`

    Args:
        state: State of the cellular automaton.

    Returns:
        Neighborhood perception.

    """
    spatial_start = state.ndim - self.num_spatial_dims - 1
    spatial_sizes = state.shape[
        spatial_start : spatial_start + self.num_spatial_dims
    ]

    pad_mode = _PADDING_TO_PAD_MODE[self.padding]
    pad_widths = (
        [(0, 0)] * spatial_start
        + [(self.radius, self.radius)] * self.num_spatial_dims
        + [(0, 0)]
    )
    padded = jnp.pad(state, pad_widths, mode=pad_mode)

    shifts = self._get_shifts()

    neighbors = []
    for shift in shifts:
        slices = [slice(None)] * spatial_start
        for dim_idx in range(self.num_spatial_dims):
            start = self.radius - shift[dim_idx]
            slices.append(slice(start, start + spatial_sizes[dim_idx]))
        slices.append(slice(None))
        neighbors.append(padded[tuple(slices)])

    if self.reduce_fn is not None:
        stacked = jnp.stack(neighbors, axis=0)
        reduced = self.reduce_fn(stacked, axis=0)
        if self.include_center:
            return jnp.concatenate([state, reduced], axis=-1)
        return reduced

    if self.include_center:
        return jnp.concatenate([state, *neighbors], axis=-1)
    return jnp.concatenate(neighbors, axis=-1)

cax.core.perceive.von_neumann_perceive

Von Neumann perceive module.

VonNeumannPerceive

Bases: NeighborhoodPerceive

Von Neumann perceive class.

This class implements perception based on the Von Neumann neighborhood. The Von Neumann neighborhood includes all cells within Manhattan distance radius of the central cell.

Source code in src/cax/core/perceive/von_neumann_perceive.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
class VonNeumannPerceive(NeighborhoodPerceive):
    """Von Neumann perceive class.

    This class implements perception based on the Von Neumann neighborhood.
    The Von Neumann neighborhood includes all cells within Manhattan distance `radius`
    of the central cell.
    """

    def __init__(
        self,
        *,
        num_spatial_dims: int,
        radius: int,
        padding: str = "CIRCULAR",
        include_center: bool = True,
        reduce_fn: Callable[..., Array] | None = None,
    ):
        """Initialize Von Neumann perceive.

        Args:
            num_spatial_dims: Number of spatial dimensions.
            radius: Manhattan distance defining the Von Neumann neighborhood extent.
            padding: Boundary condition mode. One of "CIRCULAR" (periodic/torus),
                "ZERO" (zero-padded), "REFLECT" (mirror), or "EDGE" (clamp to boundary).
            include_center: Whether to include the center cell in the output.
            reduce_fn: Optional reduction function applied over the neighbor axis. If
                None, neighbors are concatenated along the channel axis. If provided, it
                is called as `reduce_fn(stacked_neighbors, axis=0)` and the result is
                concatenated with the center (if `include_center` is True).

        """
        super().__init__(
            num_spatial_dims=num_spatial_dims,
            radius=radius,
            padding=padding,
            include_center=include_center,
            reduce_fn=reduce_fn,
        )

    @override
    def _get_shifts(self) -> list[tuple[int, ...]]:
        """Return all shifts in the Von Neumann neighborhood (excluding center).

        Returns:
            List of shift tuples with Manhattan distance <= radius, excluding the
            origin.

        """
        return [
            shift
            for shift in product(
                range(-self.radius, self.radius + 1), repeat=self.num_spatial_dims
            )
            if 0 < sum(map(abs, shift)) <= self.radius
        ]

__init__(*, num_spatial_dims, radius, padding='CIRCULAR', include_center=True, reduce_fn=None)

Initialize Von Neumann perceive.

Parameters:

Name Type Description Default
num_spatial_dims int

Number of spatial dimensions.

 required
radius int

Manhattan distance defining the Von Neumann neighborhood extent.

 required
padding str

Boundary condition mode. One of "CIRCULAR" (periodic/torus), "ZERO" (zero-padded), "REFLECT" (mirror), or "EDGE" (clamp to boundary).

 'CIRCULAR'
include_center bool

Whether to include the center cell in the output.

 True
reduce_fn Callable[..., Array] | None

Optional reduction function applied over the neighbor axis. If None, neighbors are concatenated along the channel axis. If provided, it is called as reduce_fn(stacked_neighbors, axis=0) and the result is concatenated with the center (if include_center is True).

 None
Source code in src/cax/core/perceive/von_neumann_perceive.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
def __init__(
    self,
    *,
    num_spatial_dims: int,
    radius: int,
    padding: str = "CIRCULAR",
    include_center: bool = True,
    reduce_fn: Callable[..., Array] | None = None,
):
    """Initialize Von Neumann perceive.

    Args:
        num_spatial_dims: Number of spatial dimensions.
        radius: Manhattan distance defining the Von Neumann neighborhood extent.
        padding: Boundary condition mode. One of "CIRCULAR" (periodic/torus),
            "ZERO" (zero-padded), "REFLECT" (mirror), or "EDGE" (clamp to boundary).
        include_center: Whether to include the center cell in the output.
        reduce_fn: Optional reduction function applied over the neighbor axis. If
            None, neighbors are concatenated along the channel axis. If provided, it
            is called as `reduce_fn(stacked_neighbors, axis=0)` and the result is
            concatenated with the center (if `include_center` is True).

    """
    super().__init__(
        num_spatial_dims=num_spatial_dims,
        radius=radius,
        padding=padding,
        include_center=include_center,
        reduce_fn=reduce_fn,
    )

__call__(state)

Apply neighborhood perception to the input state.

The input is assumed to have shape (..., *spatial_dims, channel_size).

When reduce_fn is None, neighbors are concatenated along the channel axis: - With include_center=True: shape (..., *spatial, channel_size * (N + 1)) - With include_center=False: shape (..., *spatial, channel_size * N) where N is the number of neighbor shifts.

When reduce_fn is provided, neighbors are stacked and reduced: - With include_center=True: shape (..., *spatial, 2 * channel_size) - With include_center=False: shape (..., *spatial, channel_size)

Parameters:

Name Type Description Default
state Array

State of the cellular automaton.

 required

Returns:

Type Description
Array

Neighborhood perception.
Source code in src/cax/core/perceive/neighborhood_perceive.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
@override
def __call__(self, state: Array) -> Array:
    """Apply neighborhood perception to the input state.

    The input is assumed to have shape `(..., *spatial_dims, channel_size)`.

    When `reduce_fn` is None, neighbors are concatenated along the channel axis:
        - With `include_center=True`: shape
          `(..., *spatial, channel_size * (N + 1))`
        - With `include_center=False`: shape `(..., *spatial, channel_size * N)`
        where N is the number of neighbor shifts.

    When `reduce_fn` is provided, neighbors are stacked and reduced:
        - With `include_center=True`: shape `(..., *spatial, 2 * channel_size)`
        - With `include_center=False`: shape `(..., *spatial, channel_size)`

    Args:
        state: State of the cellular automaton.

    Returns:
        Neighborhood perception.

    """
    spatial_start = state.ndim - self.num_spatial_dims - 1
    spatial_sizes = state.shape[
        spatial_start : spatial_start + self.num_spatial_dims
    ]

    pad_mode = _PADDING_TO_PAD_MODE[self.padding]
    pad_widths = (
        [(0, 0)] * spatial_start
        + [(self.radius, self.radius)] * self.num_spatial_dims
        + [(0, 0)]
    )
    padded = jnp.pad(state, pad_widths, mode=pad_mode)

    shifts = self._get_shifts()

    neighbors = []
    for shift in shifts:
        slices = [slice(None)] * spatial_start
        for dim_idx in range(self.num_spatial_dims):
            start = self.radius - shift[dim_idx]
            slices.append(slice(start, start + spatial_sizes[dim_idx]))
        slices.append(slice(None))
        neighbors.append(padded[tuple(slices)])

    if self.reduce_fn is not None:
        stacked = jnp.stack(neighbors, axis=0)
        reduced = self.reduce_fn(stacked, axis=0)
        if self.include_center:
            return jnp.concatenate([state, reduced], axis=-1)
        return reduced

    if self.include_center:
        return jnp.concatenate([state, *neighbors], axis=-1)
    return jnp.concatenate(neighbors, axis=-1)

cax.core.perceive.conv_perceive

Convolution perceive module.

ConvPerceive

Bases: Perceive[Array, Array]

Convolution perceive class.

Source code in src/cax/core/perceive/conv_perceive.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
class ConvPerceive(Perceive[Array, Array]):
    """Convolution perceive class."""

    def __init__(
        self,
        *,
        channel_size: int,
        perception_size: int,
        kernel_size: int | tuple[int, ...] = (3, 3),
        padding: str = "SAME",
        feature_group_count: int = 1,
        use_bias: bool = False,
        activation_fn: Callable[[Array], Array] | None = None,
        rngs: nnx.Rngs,
    ):
        """Initialize convolution perceive.

        Args:
            channel_size: Number of input channels.
            perception_size: Number of output perception features.
            kernel_size: Size of the convolutional kernel.
            padding: Padding to use.
            feature_group_count: Number of feature groups.
            use_bias: Whether to use bias in convolutional layers.
            activation_fn: Activation function to use.
            rngs: rng key.

        """
        self.conv = nnx.Conv(
            in_features=channel_size,
            out_features=perception_size,
            kernel_size=kernel_size,
            padding=padding,
            feature_group_count=feature_group_count,
            use_bias=use_bias,
            rngs=rngs,
        )
        self.activation_fn = activation_fn

    @override
    def __call__(self, state: Array) -> Array:
        """Apply convolutional perception to the input state.

        Inputs are expected to have shape `(..., *spatial_dims, channel_size)` where
        `spatial_dims` is a tuple of `num_spatial_dims` dimensions and `channel_size` is
        the number of channels. The output shape is
        `(..., *spatial_dims, perception_size)`. If `activation_fn` is provided, it is
        applied element-wise to the convolution output. If `activation_fn` is `None`,
        the convolution output is returned as is.

        Args:
            state: State of the cellular automaton.

        Returns:
            The perceived state after applying convolutional layers.

        """
        perception = self.conv(state)
        return self.activation_fn(perception) if self.activation_fn else perception

__init__(*, channel_size, perception_size, kernel_size=(3, 3), padding='SAME', feature_group_count=1, use_bias=False, activation_fn=None, rngs)

Initialize convolution perceive.

Parameters:

Name Type Description Default
channel_size int

Number of input channels.

 required
perception_size int

Number of output perception features.

 required
kernel_size int | tuple[int, ...]

Size of the convolutional kernel.

 (3, 3)
padding str

Padding to use.

 'SAME'
feature_group_count int

Number of feature groups.

 1
use_bias bool

Whether to use bias in convolutional layers.

 False
activation_fn Callable[[Array], Array] | None

Activation function to use.

 None
rngs Rngs

rng key.

 required
Source code in src/cax/core/perceive/conv_perceive.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
def __init__(
    self,
    *,
    channel_size: int,
    perception_size: int,
    kernel_size: int | tuple[int, ...] = (3, 3),
    padding: str = "SAME",
    feature_group_count: int = 1,
    use_bias: bool = False,
    activation_fn: Callable[[Array], Array] | None = None,
    rngs: nnx.Rngs,
):
    """Initialize convolution perceive.

    Args:
        channel_size: Number of input channels.
        perception_size: Number of output perception features.
        kernel_size: Size of the convolutional kernel.
        padding: Padding to use.
        feature_group_count: Number of feature groups.
        use_bias: Whether to use bias in convolutional layers.
        activation_fn: Activation function to use.
        rngs: rng key.

    """
    self.conv = nnx.Conv(
        in_features=channel_size,
        out_features=perception_size,
        kernel_size=kernel_size,
        padding=padding,
        feature_group_count=feature_group_count,
        use_bias=use_bias,
        rngs=rngs,
    )
    self.activation_fn = activation_fn

__call__(state)

Apply convolutional perception to the input state.

Inputs are expected to have shape (..., *spatial_dims, channel_size) where spatial_dims is a tuple of num_spatial_dims dimensions and channel_size is the number of channels. The output shape is (..., *spatial_dims, perception_size). If activation_fn is provided, it is applied element-wise to the convolution output. If activation_fn is None, the convolution output is returned as is.

Parameters:

Name Type Description Default
state Array

State of the cellular automaton.

 required

Returns:

Type Description
Array

The perceived state after applying convolutional layers.
Source code in src/cax/core/perceive/conv_perceive.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
@override
def __call__(self, state: Array) -> Array:
    """Apply convolutional perception to the input state.

    Inputs are expected to have shape `(..., *spatial_dims, channel_size)` where
    `spatial_dims` is a tuple of `num_spatial_dims` dimensions and `channel_size` is
    the number of channels. The output shape is
    `(..., *spatial_dims, perception_size)`. If `activation_fn` is provided, it is
    applied element-wise to the convolution output. If `activation_fn` is `None`,
    the convolution output is returned as is.

    Args:
        state: State of the cellular automaton.

    Returns:
        The perceived state after applying convolutional layers.

    """
    perception = self.conv(state)
    return self.activation_fn(perception) if self.activation_fn else perception

cax.core.perceive.kernels

Kernel utilities for perception modules.

Each function returns a small spatial kernel suitable for neighborhood aggregation or finite-difference style operations. Kernels use channel-last layout and a support of size 3 along each spatial dimension.

grad2_kernel takes a neighborhood, naming which of the surrounding cells contribute:

  • "von_neumann" --- the cells sharing a face, two per axis.
  • "moore" --- every cell in the surrounding block, diagonals included.
  • "hexagonal" --- a triangular lattice stored in a square array, whose axes are read as the basis vectors (1, 0) and (1/2, sqrt(3)/2) rather than as a Cartesian frame. Two opposite corners of the block are then not adjacent, leaving six neighbors that are all exactly one unit away. Two dimensions only.

HEX_BASIS = jnp.array([[1.0, 0.0], [0.5, math.sqrt(3.0) / 2.0]]) module-attribute

Rows are the lattice vectors a triangular lattice's array axes stand for.

identity_kernel(*, num_dims)

Create an identity kernel for the given number of dimensions.

The kernel has value 1 at the central position and 0 elsewhere.

Parameters:

Name Type Description Default
num_dims int

Number of dimensions for the kernel.

 required

Returns:

Type Description
Array

Array with shape num_dims * (3,) + (1,).
Source code in src/cax/core/perceive/kernels.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def identity_kernel(*, num_dims: int) -> Array:
    """Create an identity kernel for the given number of dimensions.

    The kernel has value 1 at the central position and 0 elsewhere.

    Args:
        num_dims: Number of dimensions for the kernel.

    Returns:
        Array with shape `num_dims * (3,) + (1,)`.

    """
    kernel = jnp.zeros(num_dims * (3,))
    center_idx = num_dims * (1,)
    kernel = kernel.at[center_idx].set(1.0)
    return jnp.expand_dims(kernel, axis=-1)

neighbors_kernel(*, num_dims)

Create a neighbors kernel for the given number of dimensions.

This kernel is 1 - identity_kernel, selecting all neighbors and excluding the center.

Parameters:

Name Type Description Default
num_dims int

Number of dimensions for the kernel.

 required

Returns:

Type Description
Array

Array with shape num_dims * (3,) + (1,).
Source code in src/cax/core/perceive/kernels.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def neighbors_kernel(*, num_dims: int) -> Array:
    """Create a neighbors kernel for the given number of dimensions.

    This kernel is `1 - identity_kernel`, selecting all neighbors and excluding the
    center.

    Args:
        num_dims: Number of dimensions for the kernel.

    Returns:
        Array with shape `num_dims * (3,) + (1,)`.

    """
    kernel = identity_kernel(num_dims=num_dims)
    return 1.0 - kernel

grad_kernel(*, num_dims, normalize=True)

Create a gradient kernel for the given number of dimensions.

Parameters:

Name Type Description Default
num_dims int

Number of dimensions for the kernel.

 required
normalize bool

Whether to L1-normalize each axis kernel.

 True

Returns:

Type Description
Array

Array with shape num_dims * (3,) + (num_dims,), one kernel per spatial axis.
Source code in src/cax/core/perceive/kernels.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
def grad_kernel(*, num_dims: int, normalize: bool = True) -> Array:
    """Create a gradient kernel for the given number of dimensions.

    Args:
        num_dims: Number of dimensions for the kernel.
        normalize: Whether to L1-normalize each axis kernel.

    Returns:
        Array with shape `num_dims * (3,) + (num_dims,)`, one kernel per spatial axis.

    """
    grad = jnp.array([-1, 0, 1])
    smooth = jnp.array([1, 2, 1])

    kernels = []
    for i in range(num_dims):
        kernel = jnp.ones([3] * num_dims)

        for j in range(num_dims):
            axis_kernel = smooth if i != j else grad
            kernel = kernel * axis_kernel.reshape(
                [-1 if k == j else 1 for k in range(num_dims)]
            )

        kernels.append(kernel)

    if normalize:
        kernels = [kernel / jnp.sum(jnp.abs(kernel)) for kernel in kernels]

    return jnp.stack(kernels, axis=-1)

grad2_kernel(*, num_dims, neighborhood='von_neumann', normalize=True)

Create a second-order (Laplacian) kernel.

Each neighbor is weighted by the inverse square of its distance from the center, and the center takes whatever makes the kernel sum to zero, so that it annihilates a constant field.

The neighborhood decides how isotropic the result is. "von_neumann" is the cheapest stencil, but its leading error favors the grid axes with four-fold symmetry --- a bias a system asked to behave the same in every direction will find and use. "moore" brings in the diagonals, which is the familiar [[1, 2, 1], [2, -12, 2], [1, 2, 1]] up to the scale normalize removes; on a quartic rotated by 45 degrees the von Neumann answer moves by 50% and this one by 20%.

"hexagonal" removes the question. Its six neighbors are all one unit away, so they are weighted equally and no direction is singled out at all. Note that an array read as a triangular lattice also has to be drawn as one --- see cax.utils.hex_to_square.

Parameters:

Name Type Description Default
num_dims int

Number of dimensions for the kernel.

 required
neighborhood Neighborhood

Which surrounding cells contribute.

 'von_neumann'
normalize bool

Whether to L1-normalize the kernel.

 True

Returns:

Type Description
Array

Array with shape num_dims * (3,) + (1,).

Raises:

Type Description
ValueError

If the neighborhood is unknown, or is hexagonal in other than two dimensions.
Source code in src/cax/core/perceive/kernels.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
def grad2_kernel(
    *, num_dims: int, neighborhood: Neighborhood = "von_neumann", normalize: bool = True
) -> Array:
    """Create a second-order (Laplacian) kernel.

    Each neighbor is weighted by the inverse square of its distance from the center, and
    the center takes whatever makes the kernel sum to zero, so that it annihilates a
    constant field.

    The neighborhood decides how isotropic the result is. `"von_neumann"` is the
    cheapest stencil, but its leading error favors the grid axes with four-fold symmetry
    --- a bias a system asked to behave the same in every direction will find and use.
    `"moore"` brings in the diagonals, which is the familiar
    `[[1, 2, 1], [2, -12, 2], [1, 2, 1]]` up to the scale `normalize` removes; on a
    quartic rotated by 45 degrees the von Neumann answer moves by 50% and this one by
    20%.

    `"hexagonal"` removes the question. Its six neighbors are all one unit away, so they
    are weighted equally and no direction is singled out at all. Note that an array read
    as a triangular lattice also has to be *drawn* as one --- see
    `cax.utils.hex_to_square`.

    Args:
        num_dims: Number of dimensions for the kernel.
        neighborhood: Which surrounding cells contribute.
        normalize: Whether to L1-normalize the kernel.

    Returns:
        Array with shape `num_dims * (3,) + (1,)`.

    Raises:
        ValueError: If the neighborhood is unknown, or is hexagonal in other than two
            dimensions.

    """
    kernel = jnp.zeros([3] * num_dims)
    for offset in _neighbor_offsets(num_dims, neighborhood):
        kernel = kernel.at[tuple(1 + step for step in offset)].set(
            1.0 / _square_distance(offset, neighborhood)
        )

    center = tuple(1 for _ in range(num_dims))
    kernel = kernel.at[center].set(-jnp.sum(kernel))

    if normalize:
        kernel = kernel / jnp.sum(jnp.abs(kernel))

    return jnp.expand_dims(kernel, axis=-1)