Source code for tbp.monty.attention.telemetry

# Copyright 2026 Thousand Brains Project
#
# Copyright may exist in Contributors' modifications
# and/or contributions to the work.
#
# Use of this source code is governed by the MIT
# license that can be found in the LICENSE file or at
# https://opensource.org/licenses/MIT.
from __future__ import annotations

from typing import Protocol

from tbp.monty.attention.voxel_grid import VoxelGrid
from tbp.monty.memento import Memento

__all__ = [
    "AttentionSystemTelemetry",
    "DefaultAttentionSystemTelemetry",
    "NoopAttentionSystemTelemetry",
]


[docs]class AttentionSystemTelemetry(Protocol):
[docs]    def reset(self) -> None: ...

[docs]    def proposed_grid(self, grid: VoxelGrid) -> None: ...

[docs]    def grid(self, grid: VoxelGrid) -> None: ...

[docs]    def state_dict(self) -> Memento: ...


[docs]class NoopAttentionSystemTelemetry(AttentionSystemTelemetry):
[docs]    def reset(self) -> None:
        pass

[docs]    def proposed_grid(self, grid: VoxelGrid) -> None:
        pass

[docs]    def grid(self, grid: VoxelGrid) -> None:
        pass

[docs]    def state_dict(self) -> Memento:
        return dict(grids=[], proposed_grids=[])


[docs]class DefaultAttentionSystemTelemetry(AttentionSystemTelemetry):
[docs]    def __init__(self) -> None:
        self._grids: list[VoxelGrid] = []
        self._proposed_grids: list[VoxelGrid] = []

[docs]    def reset(self) -> None:
        self._grids = []
        self._proposed_grids = []

[docs]    def grid(self, grid: VoxelGrid) -> None:
        self._grids.append(grid)

[docs]    def proposed_grid(self, grid: VoxelGrid) -> None:
        self._proposed_grids.append(grid)

[docs]    def state_dict(self) -> Memento:
        return dict(
            grids=list(self._grids),
            proposed_grids=list(self._proposed_grids),
        )