| """MLX dual-timescale matrix state with denoise and commit lifetimes.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import mlx.core as mx |
| from mlx import nn |
|
|
| from .mlx_gdn2_memory import GDN2Memory |
|
|
|
|
| @dataclass(frozen=True) |
| class GDN2TrajectoryState: |
| cells: mx.array |
| row: mx.array |
| persistent: mx.array |
| seen: mx.array |
|
|
|
|
| def clear_refill(self, lengths: mx.array, head: mx.array) -> "GDN2TrajectoryState": |
| batch, canvas = self.seen.shape |
| if lengths.shape != (batch,) or head.shape != (batch,): |
| raise ValueError("Commit lengths and canvas heads must be per row.") |
| physical = mx.arange(canvas)[None, :] |
| overwritten = ((physical - head[:, None]) % canvas) < lengths[:, None] |
| return GDN2TrajectoryState( |
| mx.where(overwritten[..., None, None, None], 0.0, self.cells), |
| self.row, self.persistent, self.seen & ~overwritten, |
| ) |
|
|
|
|
| class GDN2TrajectoryMemory(nn.Module): |
| def __init__(self, width: int, *, probes: int = 4, |
| working_heads: int = 16, working_key: int = 64, |
| working_value: int = 64, persistent_heads: int = 16, |
| persistent_key: int = 128, persistent_value: int = 128, |
| persistent_observation_dim: int | None = None) -> None: |
| super().__init__() |
| if probes <= 0: |
| raise ValueError("Probe count must be positive.") |
| self.probes = probes |
| self.cell = GDN2Memory(width, working_heads, working_key, working_value) |
| self.row = GDN2Memory(width, working_heads, working_key, working_value) |
| self.persistent = GDN2Memory(width, persistent_heads, persistent_key, persistent_value, |
| observation_dim=persistent_observation_dim) |
| self.probe_embed = nn.Embedding(probes, width) |
| self.probe_embed.weight = mx.random.normal((probes, width)) * 0.02 |
|
|
| def empty(self, batch: int, canvas: int) -> GDN2TrajectoryState: |
| return GDN2TrajectoryState( |
| self.cell.empty(batch, canvas), self.row.empty(batch), |
| self.persistent.empty(batch), mx.zeros((batch, canvas), mx.bool_), |
| ) |
|
|
| def read(self, state: GDN2TrajectoryState, query: mx.array) -> mx.array: |
| result = (self.cell.read(state.cells, query) |
| + self.row.read_shared(state.row, query) |
| + self.persistent.read_shared(state.persistent, query)) |
| return mx.where(state.seen[..., None], result, 0.0) |
|
|
| def observe(self, state: GDN2TrajectoryState, observation: mx.array, |
| live: mx.array, head: mx.array) -> GDN2TrajectoryState: |
| batch, canvas, width = observation.shape |
| if live.shape != (batch, canvas) or head.shape != (batch,): |
| raise ValueError("Observation mask and head have incorrect shapes.") |
| cells = self.cell.transition(state.cells, observation, live) |
| seen = state.seen | live |
| logical_idx = (head[:, None] + mx.arange(canvas)[None, :]) % canvas |
| logical = mx.take_along_axis(observation, |
| mx.broadcast_to(logical_idx[..., None], observation.shape), axis=1) |
| logical_live = mx.take_along_axis(live, mx.stop_gradient(logical_idx), axis=1) |
| row_state = state.row |
| for probe in range(self.probes): |
| lo = canvas * probe // self.probes |
| hi = canvas * (probe + 1) // self.probes |
| selected = logical_live[:, lo:hi] |
| count = mx.sum(selected.astype(mx.float32), axis=1, keepdims=True) |
| pooled = mx.sum(logical[:, lo:hi].astype(mx.float32) |
| * selected[..., None], axis=1) |
| pooled = (pooled / mx.maximum(count, 1.0)).astype(observation.dtype) |
| pooled = pooled + self.probe_embed.weight[probe].astype(observation.dtype) |
| row_state = self.row.transition(row_state, pooled, count[:, 0] > 0) |
| return GDN2TrajectoryState(cells, row_state, state.persistent, seen) |
|
|
|
|
| __all__ = ["GDN2TrajectoryMemory", "GDN2TrajectoryState"] |
|
|
