File size: 4,089 Bytes
e4f7326 | 1 2 3 4 5 6 7 8 9 10 11 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 80 81 82 83 84 85 86 87 88 | """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"]
|