""" Online (TBPTT) training: Fixed-ZOH HiPPO + Continuous-time Associative Memory on the associative recall task (A/B tokens + WRITE). This script trains a fixed-ZOH HiPPO model with a continuous-time associative memory module for associative recall. The model precomputes one ZOH discretization, keeps HiPPO state S and memory bank C as separate states, applies exact memory updates each step, and trains with TBPTT using stop_gradient at chunk boundaries. """ __date__ = "January 2026" import numpy as np import matplotlib.pyplot as plt import jax import jax.numpy as jnp import jax.random as jr from jax import jit, value_and_grad from jax.lax import scan import flax.linen as nn import optax from tqdm import tqdm from mpm import get_system_params from mpm.models import AssocMemHiPPO, legendre_orthonormal_basis01 from mpm.tasks.associative_memory import make_associative_recall_task_tokens # ---------------------------- # TBPTT chunk padding # ---------------------------- def pad_chunk(X, Y, t0: int, L: int): T = X.shape[0] t1 = min(t0 + L, T) ell = t1 - t0 Xc = jnp.zeros((L, X.shape[1]), dtype=X.dtype) Yc = jnp.zeros((L, Y.shape[1]), dtype=Y.dtype) mask = jnp.zeros((L,), dtype=X.dtype) Xc = Xc.at[:ell].set(X[t0:t1]) Yc = Yc.at[:ell].set(Y[t0:t1]) mask = mask.at[:ell].set(1.0) return Xc, Yc, mask # ---------------------------- # Training (TBPTT) # ---------------------------- def main(): # ---------------------------- # Config # ---------------------------- seed = 42 key = jr.PRNGKey(seed) # task episode_len = 12 # must be even num_episodes = 10000 num_tokens = 12 # online training steps (episodes sampled) num_steps = 3500 # TBPTT tbptt_steps = episode_len # token dims / model dims d_in = 24 d_model = 32 write_hidden = 256 out_hidden = 256 # HiPPO n_hippo = 32 measure = "legt" base_timescale = 2.0 # Associative memory truncation n_assoc = 32 # opt lr = 1e-3 wd = 1e-4 # ---------------------------- # Dataset # ---------------------------- X_np, Y_np, token_table, ids_np, meta_np = make_associative_recall_task_tokens( episode_len=episode_len, num_episodes=num_episodes, d_in=d_in, num_tokens=num_tokens, seed=seed, ) X_all = jnp.array(X_np, dtype=jnp.float32) # (N,T,d_in) Y_all = jnp.array(Y_np, dtype=jnp.float32) ids_all = np.array(ids_np) # ---------------------------- # HiPPO system + fixed ZOH discretization (dt=1) # ---------------------------- (A, b), _, _ = get_system_params(measure, n_hippo) # scale time A = A / base_timescale b = b / base_timescale A = jnp.array(A, dtype=jnp.float32) b = jnp.array(b, dtype=jnp.float32) I = jnp.eye(n_hippo, dtype=jnp.float32) A_d = jax.scipy.linalg.expm(A) # dt = 1 # b_d = \int_0^1 exp(A \tau) b d\tau = A^{-1}(A_d - I)b (use solve for stability) rhs = (A_d - I) @ b b_d = jnp.linalg.solve(A, rhs) # ---------------------------- # Init model + optimizer # ---------------------------- model = AssocMemHiPPO( d_in=d_in, d_model=d_model, n_hippo=n_hippo, n_assoc=n_assoc, A_d=A_d, b_d=b_d, write_hidden=write_hidden, out_hidden=out_hidden, ) key, k1 = jr.split(key) x0 = jnp.zeros((d_in,), dtype=jnp.float32) S0 = jnp.zeros((d_model, n_hippo), dtype=jnp.float32) C0 = jnp.zeros((d_model, n_assoc), dtype=jnp.float32) params = model.init(k1, x0, S0, C0)["params"] optimizer = optax.adamw(learning_rate=lr, weight_decay=wd) opt_state = optimizer.init(params) # ---------------------------- # Chunk rollout + TBPTT update # ---------------------------- def chunk_loss_and_state(params, S_init, C_init, Xc, Yc, mask): """ One TBPTT chunk. S_init: (d_model,n_hippo) C_init: (d_model,n_assoc) Xc,Yc: (L,d_in) padded mask: (L,) in {0,1} """ def step_fn(carry, inputs): S, C = carry x_t, y_t, m_t = inputs S_next, C_next, yhat, aux = model.apply({"params": params}, x_t, S, C) mse_t = jnp.mean((yhat - y_t) ** 2) return (S_next, C_next), (m_t * mse_t, m_t, yhat, aux["g_write"], aux["g_out"]) (S_final, C_final), (mse_masked, m_sum, yhat_ts, g_write_ts, g_out_ts) = scan( step_fn, (S_init, C_init), (Xc, Yc, mask), ) denom = jnp.maximum(jnp.sum(m_sum), 1.0) loss = jnp.sum(mse_masked) / denom aux = (yhat_ts, g_write_ts, g_out_ts) return loss, (S_final, C_final), aux @jit def tbptt_chunk_train_step(params, opt_state, S_init, C_init, Xc, Yc, mask): def loss_fn(p): loss, (S_final, C_final), aux = chunk_loss_and_state(p, S_init, C_init, Xc, Yc, mask) return loss, (S_final, C_final, aux) (loss, (S_final, C_final, aux)), grads = value_and_grad(loss_fn, has_aux=True)(params) updates, opt_state = optimizer.update(grads, opt_state, params) params = optax.apply_updates(params, updates) return params, opt_state, S_final, C_final, loss, aux # ---------------------------- # Online training loop (episode-by-episode) with TBPTT # ---------------------------- losses = [] key = jr.PRNGKey(seed + 1) N = X_all.shape[0] pbar = tqdm(range(num_steps)) smooth_loss = None for step in pbar: key, sub = jr.split(key) idx = jr.randint(sub, (), 0, N) X_ep = X_all[idx] # (T,d_in) Y_ep = Y_all[idx] # (T,d_in) # Reset states per episode S = jnp.zeros((d_model, n_hippo), dtype=jnp.float32) C = jnp.zeros((d_model, n_assoc), dtype=jnp.float32) # TBPTT over chunks T = X_ep.shape[0] n_chunks = (T + tbptt_steps - 1) // tbptt_steps loss_ep = 0.0 for c in range(n_chunks): t0 = c * tbptt_steps Xc, Yc, mask = pad_chunk(X_ep, Y_ep, t0=t0, L=tbptt_steps) params, opt_state, S_end, C_end, loss_chunk, _ = tbptt_chunk_train_step( params, opt_state, S, C, Xc, Yc, mask ) loss_ep = loss_ep + loss_chunk # stop gradients across TBPTT boundaries S = jax.lax.stop_gradient(S_end) C = jax.lax.stop_gradient(C_end) loss_ep = loss_ep / n_chunks losses.append(float(loss_ep)) if smooth_loss is None: smooth_loss = losses[-1] else: smooth_loss = 0.99 * smooth_loss + 0.01 * losses[-1] pbar.set_description(f"loss: {smooth_loss:.6f}") # ---------------------------- # Evaluate on one held-out episode # ---------------------------- X_eval = jnp.concatenate([X_all[-2], X_all[-1]], 0) Y_eval = jnp.concatenate([Y_all[-2], Y_all[-1]], 0) def full_rollout_collect(params, X, Y): S = jnp.zeros((d_model, n_hippo), dtype=jnp.float32) C = jnp.zeros((d_model, n_assoc), dtype=jnp.float32) def step_fn(carry, inputs): S, C = carry x_t, y_t = inputs S_next, C_next, yhat, aux = model.apply({"params": params}, x_t, S, C) loss_t = jnp.mean((yhat - y_t) ** 2) return (S_next, C_next), (loss_t, yhat, aux["g_write"], aux["g_out"]) _, (loss_ts, yhat_ts, g_write_ts, g_out_ts) = scan(step_fn, (S, C), (X, Y)) return jnp.mean(loss_ts), yhat_ts, g_write_ts, g_out_ts loss_eval, preds_ts, g_write_ts, g_out_ts = full_rollout_collect(params, X_eval, Y_eval) preds = np.array(preds_ts) Yp = np.array(Y_eval) g_write_np = np.array(g_write_ts) g_out_np = np.array(g_out_ts) def cos_sim(a, b, eps=1e-8): na = np.linalg.norm(a, axis=-1) nb = np.linalg.norm(b, axis=-1) return np.sum(a * b, axis=-1) / (na * nb + eps) cos = cos_sim(preds, Yp) t = np.arange(len(Y_eval)) fig, ax = plt.subplots(4, 1, figsize=(10, 9)) ax[0].set_title(f"Associative recall (eval loss={float(loss_eval):.4g})") ax[0].plot(np.arange(len(losses)), losses) ax[0].set_ylabel("train loss") for k in [0, 1, 2]: ax[1].plot(t, Yp[:, k], label=f"Y dim{k}", linewidth=2) ax[1].plot(t, preds[:, k], linestyle="--", label=f"pred dim{k}", alpha=0.8) ax[1].legend(loc="upper left", ncol=3) ax[1].set_ylabel("dims") ax[2].plot(t, cos) ax[2].set_ylabel("cos(pred,target)") ax[2].set_ylim(-0.05, 1.05) ax[3].plot(t, g_write_np, label="g write") ax[3].plot(t, g_out_np, label="g out") ax[3].set_ylabel("g") ax[3].legend(loc="best") ax[3].set_xlabel("timestep") ax[3].set_ylim(0, None) plt.tight_layout() plt.savefig("temp.png") plt.close('all') make_temp2_plot( model=model, params=params, X_all=np.array(X_all), ids_all=ids_all, token_table=token_table, num_tokens=num_tokens, n_holdout=200, seed=0, x_grid_G=250, ) print("Saved temp2.png") def _cos_sim(a, b, eps=1e-8): # a: (..., d), b: (..., d) or (d,) a = np.asarray(a) b = np.asarray(b) na = np.linalg.norm(a, axis=-1) nb = np.linalg.norm(b, axis=-1) return np.sum(a * b, axis=-1) / (na * nb + eps) def kernel_overlap_legendre01(x1, x2, n_assoc): """ Normalized kernel overlap: k(x1,x2) / sqrt(k(x1,x1)k(x2,x2)), where k(x,z) = phi(x)^T phi(z) for orthonormal basis. """ phi1 = np.array(legendre_orthonormal_basis01(jnp.asarray(x1), n_assoc)) phi2 = np.array(legendre_orthonormal_basis01(jnp.asarray(x2), n_assoc)) k12 = float(np.dot(phi1, phi2)) k11 = float(np.dot(phi1, phi1)) k22 = float(np.dot(phi2, phi2)) return k12 / (np.sqrt(k11 * k22) + 1e-8) def rollout_collect_episode(model, params, X_ep, ids_ep): """ Runs one episode and collects aux + memory states. Returns dict with: x_key: (T,) x_query: (T,) g_write: (T,) g_out: (T,) C_ts: (T, d_model, n_assoc) post-step C """ # infer shapes from model config d_model = model.d_model n_hippo = model.n_hippo n_assoc = model.n_assoc S0 = jnp.zeros((d_model, n_hippo), dtype=jnp.float32) C0 = jnp.zeros((d_model, n_assoc), dtype=jnp.float32) def step_fn(carry, x_t): S, C = carry S_next, C_next, y_hat, aux = model.apply({"params": params}, x_t, S, C) return (S_next, C_next), (aux["x_key"], aux["x_query"], aux["g_write"], aux["g_out"], C_next) (Sf, Cf), (x_key_ts, x_query_ts, g_write_ts, g_out_ts, C_ts) = jax.lax.scan(step_fn, (S0, C0), X_ep) return { "x_key": np.array(x_key_ts), "x_query": np.array(x_query_ts), "g_write": np.array(g_write_ts), "g_out": np.array(g_out_ts), "C_ts": np.array(C_ts), "ids": np.array(ids_ep), } def find_target_write_time(ids_ep, num_tokens): """ For this task structure: - WRITE token id is 2*num_tokens - final A token is at time T-2 (even index) - find last earlier occurrence of that A token among even indices < T-2 - write time is the B immediately after that earlier A: t_write = tA + 1 Returns: (t_write, a_id, b_vocab_id) where b_vocab_id is in [num_tokens, 2*num_tokens-1] """ WRITE_ID = 2 * num_tokens T = len(ids_ep) assert ids_ep[-1] == WRITE_ID, "expected WRITE at final position" a_id = int(ids_ep[T - 2]) # final A token id in [0,num_tokens-1] A_times = np.arange(0, T - 1, 2) # even positions excluding final WRITE A_ids = ids_ep[A_times] # find last earlier occurrence among A_times excluding last A matches = np.where(A_ids[:-1] == a_id)[0] if len(matches) == 0: return None # should not happen if dataset generator enforced it j_last = int(matches[-1]) tA = int(A_times[j_last]) t_write = tA + 1 b_vocab = int(ids_ep[t_write]) return t_write, a_id, b_vocab def decode_memory_curve(C_bank, out_kernel, out_bias, x_grid): """ C_bank: (d_model, n_assoc) associative memory banks at some time out_kernel: (d_model, d_in) out_bias: (d_in,) x_grid: (G,) in [0,1] Returns: Y_mem: (G, d_in) decoded token-space vector as a function of x """ d_model, n_assoc = C_bank.shape G = len(x_grid) # build Phi(x): (G, n_assoc) Phi = np.array(legendre_orthonormal_basis01(jnp.asarray(x_grid), n_assoc)) # (G, n_assoc) # r(x) = C_bank @ phi(x): (d_model,) for each x => (G, d_model) R = Phi @ C_bank.T # (G, d_model) # out_proj: y = R @ W + b => (G, d_in) Y = R @ out_kernel + out_bias[None, :] return Y # ---------------------------- # Diagnostics on holdout episodes # ---------------------------- def make_temp2_plot( model, params, X_all, ids_all, token_table, num_tokens: int, n_holdout: int = 200, seed: int | None = None, x_grid_G: int = 200, ): rng = np.random.default_rng(seed) N, T, d_in = X_all.shape # Choose holdout indices: here we just sample from the *end* chunk of the dataset. # If you have a true split, replace this block. holdout_pool = np.arange(int(0.8 * N), N) sel = rng.choice(holdout_pool, size=min(n_holdout, len(holdout_pool)), replace=False) # Collect per-episode (read_addr, write_addr, a_id, b_vocab) read_addrs = [] write_addrs = [] a_ids = [] b_vocabs = [] per_ep_rollouts = [] # keep a few for plot 3 selection for idx in sel: X_ep = X_all[idx] ids_ep = ids_all[idx] # rollout roll = rollout_collect_episode(model, params, X_ep, ids_ep) per_ep_rollouts.append(roll) # identify write time corresponding to the retrieved association out = find_target_write_time(roll["ids"], num_tokens=num_tokens) if out is None: continue t_write, a_id, b_vocab = out # read address at WRITE time t_read = T - 1 x_read = float(roll["x_query"][t_read]) # write address at the *write time* (where the relevant B appears) x_write = float(roll["x_key"][t_write]) read_addrs.append(x_read) write_addrs.append(x_write) a_ids.append(a_id) b_vocabs.append(b_vocab) read_addrs = np.array(read_addrs) write_addrs = np.array(write_addrs) a_ids = np.array(a_ids, dtype=np.int32) b_vocabs = np.array(b_vocabs, dtype=np.int32) # ---------------------------- # Plot 1: scatter read vs write addresses, colored by A token # ---------------------------- fig, axarr = plt.subplots(3, 1, figsize=(10, 14)) ax1, ax2, ax3 = axarr cmap = plt.get_cmap("tab10") colors = [cmap(i % 10) for i in range(num_tokens)] for a in range(num_tokens): m = (a_ids == a) if np.any(m): ax1.scatter(write_addrs[m], read_addrs[m], s=20, alpha=0.8, color=colors[a], label=f"A{a}") ax1.set_title("1) Read vs write address for the retrieved association (colored by A token)") ax1.set_xlabel("write address x_key at relevant [a;b] time") ax1.set_ylabel("read address x_query at [a;WRITE] time") ax1.set_xlim(-0.02, 1.02) ax1.set_ylim(-0.02, 1.02) ax1.plot([0,1], [0,1], c='k', alpha=0.5, ls='--') ax1.grid(True, alpha=0.2) ax1.legend(loc="upper right", ncol=2, fontsize=9) # # ---------------------------- # Plot 2: kernel overlap between mean write addresses per A token + null band # ---------------------------- mean_write_addr = np.full((num_tokens,), np.nan, dtype=np.float32) for a in range(num_tokens): m = (a_ids == a) if np.any(m): mean_write_addr[a] = float(np.mean(write_addrs[m])) # Pair overlaps for A tokens pair_overlaps = [] pair_labels = [] for i in range(num_tokens): for j in range(i + 1, num_tokens): if np.isfinite(mean_write_addr[i]) and np.isfinite(mean_write_addr[j]): ov = kernel_overlap_legendre01(mean_write_addr[i], mean_write_addr[j], n_assoc=model.n_assoc) pair_overlaps.append(ov) pair_labels.append((i, j)) pair_overlaps = np.array(pair_overlaps, dtype=np.float32) # ---------------------------- # Null KDE of max |overlap| across num_tokens random addresses # ---------------------------- def max_abs_pair_overlap_from_addresses(xs: jnp.ndarray, n_assoc: int) -> jnp.ndarray: """ xs: (M,) addresses in [0,1] returns: scalar max_{i / (||phi(xi)|| ||phi(xj)||) | """ Phi = legendre_orthonormal_basis01(xs, n_assoc) # (M, n_assoc) Phi = Phi / (jnp.linalg.norm(Phi, axis=1, keepdims=True) + 1e-8) G = Phi @ Phi.T # (M, M), diag ~ 1 M = xs.shape[0] # upper triangle mask without diag mask = jnp.triu(jnp.ones((M, M), dtype=G.dtype), k=1) return jnp.max(jnp.abs(G) * mask) max_abs_pair_overlap_from_addresses_jit = jax.jit( max_abs_pair_overlap_from_addresses, static_argnames=("n_assoc",) ) def null_max_abs_overlap_distribution( key: jax.Array, num_tokens: int, n_assoc: int, n_trials: int, ) -> jnp.ndarray: """ Draw n_trials sets of num_tokens addresses ~ Uniform[0,1], return distribution of max abs pair overlap. Shape (n_trials,). """ # sample all addresses at once: (n_trials, num_tokens) U = jr.uniform(key, shape=(n_trials, num_tokens), minval=0.0, maxval=1.0) # vectorize across trials f = lambda xs: max_abs_pair_overlap_from_addresses_jit(xs, n_assoc=n_assoc) return jax.vmap(f)(U) # (n_trials,) # ---- observed statistic from mean write addresses (computed earlier) # mean_write_addr: (num_tokens,) with NaNs possible if token missing valid = np.isfinite(mean_write_addr) xs_obs = mean_write_addr[valid].astype(np.float32) m_eff = xs_obs.shape[0] if m_eff < 2: ax2.set_title("2) Not enough observed tokens with write addresses to compute overlaps") else: # observed max abs overlap T_obs = float( max_abs_pair_overlap_from_addresses_jit(jnp.asarray(xs_obs), n_assoc=int(model.n_assoc)) ) # null distribution n_trials = 1000 # increase if you want smoother KDE key_null = jr.PRNGKey(123) # or fold in your main key/seed T_null = np.array( null_max_abs_overlap_distribution( key_null, num_tokens=m_eff, n_assoc=int(model.n_assoc), n_trials=n_trials ), dtype=np.float32 ) # ---- KDE plot (SciPy if available; else histogram) ax2.cla() from scipy.stats import gaussian_kde kde = gaussian_kde(T_null) xs = np.linspace(0.0, max(1e-3, float(np.max(T_null)) * 1.05), 400) ys = kde(xs) ax2.plot(xs, ys, linewidth=2.5, label=f"Null KDE of max |overlap| (n={n_trials})") ax2.fill_between(xs, 0.0, ys, alpha=0.25) ax2.axvline(T_obs, linewidth=3.0, linestyle="--", label=f"Observed max |overlap| = {T_obs:.3f}") # useful summary stats p = float(np.mean(T_null >= T_obs)) q95 = float(np.quantile(T_null, 0.95)) q99 = float(np.quantile(T_null, 0.99)) ax2.set_title( "2) Max pairwise |kernel overlap| among A-token mean write addresses\n" f"null >= observed fraction p~={p:.3f} (null 95%={q95:.3f}, 99%={q99:.3f})" ) ax2.set_xlabel("max_{i best_score: best_score = score best = roll best_info = (t_write, a_id, b_vocab) if best is None: ax3.set_title("3) (Could not find valid episode for before/after write plot)") else: t_write, a_id, b_vocab = best_info # We have C_ts as post-step. Approximate: # - "before" write: C just before processing t_write -> use C_ts[t_write-1] # - "after" write: C after processing t_write -> use C_ts[t_write] # If t_write==0 (shouldn't happen), fallback to t_write. C_before = best["C_ts"][max(t_write - 1, 0)] C_after = best["C_ts"][t_write] # out_proj params W_out = np.array(params["out_proj"]["kernel"]) # (d_model, d_in) b_out = np.array(params["out_proj"]["bias"]) # (d_in,) x_grid = np.linspace(0.0, 1.0, x_grid_G, dtype=np.float32) Y_before = decode_memory_curve(C_before, W_out, b_out, x_grid) # (G, d_in) Y_after = decode_memory_curve(C_after, W_out, b_out, x_grid) # (G, d_in) # B token vectors in token_table: vocab ids num_tokens..2*num_tokens-1 # We'll plot cos( memory(x), token_b ) as a function of x. for b_id in range(num_tokens): vocab_id = num_tokens + b_id v = token_table[vocab_id] # (d_in,) c_before = _cos_sim(Y_before, v[None, :]) c_after = _cos_sim(Y_after, v[None, :]) # Plot after as solid, before as dashed (light) ax3.plot(x_grid, c_after, linewidth=2.0, alpha=0.85, label=f"B{b_id}" if num_tokens <= 10 else None) ax3.plot(x_grid, c_before, linewidth=1.2, alpha=0.35, linestyle="--") # Mark the write address at t_write (where b_vocab was observed) for reference xw = float(best["x_key"][t_write]) ax3.axvline(xw, linestyle=":", linewidth=2.0, alpha=0.8) written_b = int(b_vocab - num_tokens) # in 0..num_tokens-1 ax3.set_title( "3) Before (dashed) vs After (solid) write: cos(memory(x), B-token)\n" f"selected episode: g_write@t_write={best_score:.3f}, wrote B{written_b} at x_key={xw:.3f}" ) ax3.set_xlabel("x in [0,1] (OP address)") ax3.set_ylabel("cosine similarity") # ax3.set_ylim(-0.05, 1.05) ax3.grid(True, alpha=0.2) if num_tokens <= 10: ax3.legend(loc="upper right", ncol=2, fontsize=9) plt.tight_layout() plt.savefig("temp2.png", dpi=200) plt.close(fig) return mean_write_addr @jax.jit(static_argnames=("n_assoc",)) def null_kernel_overlaps_legendre01(u1: jnp.ndarray, u2: jnp.ndarray, n_assoc: int) -> jnp.ndarray: """ Vectorized (Option A) null kernel overlaps for orthonormal Legendre basis on [0,1]. Computes normalized overlaps: ov[i] = / (||phi(u1[i])|| * ||phi(u2[i])||) Args: u1, u2: shape (N,) in [0,1] n_assoc: truncation/order (python int is fine) Returns: overlaps: shape (N,) float32 """ u1 = jnp.asarray(u1) u2 = jnp.asarray(u2) # Phi1, Phi2: (N, n_assoc) Phi1 = legendre_orthonormal_basis01(u1, n_assoc) Phi2 = legendre_orthonormal_basis01(u2, n_assoc) # Dot products and norms: (N,) dot12 = jnp.sum(Phi1 * Phi2, axis=-1) n1 = jnp.sqrt(jnp.sum(Phi1 * Phi1, axis=-1) + 1e-8) n2 = jnp.sqrt(jnp.sum(Phi2 * Phi2, axis=-1) + 1e-8) return (dot12 / (n1 * n2)).astype(jnp.float32) @jax.jit(static_argnames=("n_assoc", "grid_n")) def kernel_overlap_grid_legendre01( n_assoc: int, grid_n: int, ) -> jnp.ndarray: """ Compute pairwise normalized kernel overlap on a uniform grid in [0,1]. Args: n_assoc: truncation/order of orthonormal Legendre basis grid_n: number of grid points (includes endpoints 0 and 1) Returns: K: (grid_n, grid_n) array with K[i,j] = / (||phi(x_i)|| ||phi(x_j)||) """ # Uniform grid including endpoints x = jnp.linspace(0.0, 1.0, grid_n) # Basis evaluations: Phi[i,k] = p_k(x_i) Phi = legendre_orthonormal_basis01(x, n_assoc) # (grid_n, n_assoc) # Normalize basis vectors Phi_norm = Phi / (jnp.linalg.norm(Phi, axis=1, keepdims=True) + 1e-8) # Pairwise normalized overlaps K = Phi_norm @ Phi_norm.T # (grid_n, grid_n) return K.astype(jnp.float32) ARR = [ 0.45298209190368655, 0.21947033051401377, 0.5032657034256879, 0.2627926245331764, 0.5524154046307439, 0.4039508490001454, 0.6963000237941742, 0.3547567844390869, 0.7415585688182286, 0.3075451672077179, 0.6019835743037137, 0.6492930816279517, ] # Empirical mean write & read addresses def make_kernel_plot(mean_write_addr=ARR): grid = kernel_overlap_grid_legendre01(32, 256) grid = np.clip(np.array(grid), -1, 1) print(np.min(grid), np.max(grid)) xx, yy = np.meshgrid(mean_write_addr, mean_write_addr, indexing='xy') plt.scatter(xx, yy, marker='+', c='k', s=9.0, alpha=0.7) plt.imshow(grid, vmin=-1, vmax=1, extent=(0,1,0,1), cmap='bwr', origin='lower') plt.xlabel(r'Address $x_{query}$') plt.ylabel(r'Address $x_{key}$') plt.title(r'OP Memory Kernel $K(x_{key}, x_{query})$ ($n_{assoc}=32$)') plt.colorbar() plt.savefig('op_memory_kernel.png') plt.savefig('op_memory_kernel.pdf') plt.close('all') if __name__ == "__main__": mean_write_addr = main() # make_kernel_plot()