Transformer predictive coding networks

This notebook builds the smallest possible GPT-style Transformer in JAX/Equinox and trains it twice with ordinary backpropagation and with PyHGF using the hybrid approach described here.

The two learners are then put side by side and both generate text. The pipeline is deliberately minimal and self-contained: a char-level tokenizer, a tiny GPT (2 layers, 64-dim, 4 heads), short training loops, and autoregressive text generation.

The PyHGF machinery assembles a mixed pipeline: a sequence of parts where each slot holds either a frozen calculation (normalisation, the attention mixing, things with nothing to learn, which only translate errors with hand-derived formulas) or a learning PyHGF network. Walked forward, the pipeline predicts. Walked backward, each part learns locally and hands the error at its input to the part behind it.

import time

import equinox as eqx
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import optax
import seaborn as sns
from jax import random

from pyhgf.model import (
    DeepNetworkAdapter,
    FusedPipeline,
    from_embedding,
    from_feedforward,
    from_linear,
    hybrid_from_gpt,
)

plt.rcParams["figure.constrained_layout.use"] = True

Data and tokenizer

We use a character-level tokenizer: the vocabulary is simply the set of unique characters in the corpus. stoi/itos map characters to integer ids and back. This needs no external library and keeps the vocabulary tiny (a few dozen tokens), which is ideal for a minimal model.

The corpus is a short thematic text repeated a few times so the tiny model has enough sequence windows to learn a reproducible pattern.

corpus = (
    "predictive coding networks learn by minimizing prediction errors. "
    "each layer predicts the activity of the next, and updates its weights locally. "
) * 30

chars = sorted(set(corpus))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
vocab_size = len(chars)

encode = lambda s: jnp.array([stoi[c] for c in s], dtype=jnp.int32)
decode = lambda ids: "".join(itos[int(i)] for i in ids)

data = encode(corpus)
print(
    f"vocab_size = {vocab_size}, corpus length = {len(data)} tokens, vocabulary: {''.join(chars)!r}"
)

vocab_size = 27, corpus length = 4350 tokens, vocabulary: ' ,.abcdefghiklmnoprstuvwxyz'

Batching

Training examples are random windows of length T (the context length). For a window starting at i, the input is data[i:i+T] and the target is the same window shifted by one (data[i+1:i+T+1]), i.e. predict the next token at every position.

T = 32  # context length


def get_batch(key, batch_size):
    """Draw ``batch_size`` windows of ``T`` characters and their next characters."""
    ix = random.randint(key, (batch_size,), 0, len(data) - T - 1)
    windows = ix[:, None] + jnp.arange(T)  # one gather, no Python loop
    return data[windows], data[windows + 1]


xb, yb = get_batch(random.key(0), batch_size=4)
print("x:", xb.shape, "y:", yb.shape)
print("example input :", decode(xb[0]))
print("example target:", decode(yb[0]))

x: (4, 32) y: (4, 32)
example input : orks learn by minimizing predict
example target: rks learn by minimizing predicti

The architecture

A GPT is a feedforward stack (no recurrence): token+position embeddings, then L identical blocks, then a final norm and an unembedding head. Each block is two sub-layers, each wrapped in pre-norm + residual:

x  ->  x + Attention(LayerNorm(x))   # mixes information ACROSS positions
x  ->  x + FeedForward(LayerNorm(x)) # acts on each position INDEPENDENTLY

Following the Equinox convention, every module operates on a single sequence x of shape (T, D); we jax.vmap over the batch only at the loss. Per-token linear layers are applied with jax.vmap over the T axis.

Multi-head causal self-attention

The only operation that moves information between token positions. The W_{Q,K,V,O} are learnable; the mixing matrix softmax(QKᵀ/√d) is computed from the data (no weights). The causal mask forbids attending to future positions.

class MultiHeadAttention(eqx.Module):
    """Causal multi-head self-attention over one sequence of shape ``(T, D)``."""

    wqkv: eqx.nn.Linear
    wo: eqx.nn.Linear
    n_heads: int = eqx.field(static=True)
    head_dim: int = eqx.field(static=True)

    def __init__(self, dim, n_heads, key):
        k1, k2 = random.split(key)
        # The query, key, and value projections read the same input, so they
        # are one fused matrix: the output stacks the three streams on the
        # feature axis, [q | k | v], computed as a single matrix product.
        self.wqkv = eqx.nn.Linear(dim, 3 * dim, use_bias=False, key=k1)
        self.wo = eqx.nn.Linear(dim, dim, use_bias=False, key=k2)
        self.n_heads = n_heads
        self.head_dim = dim // n_heads

    def __call__(self, x):
        """Mix information across the positions of ``x`` of shape ``(T, D)``."""
        seq_len, dim = x.shape
        q, k, v = jnp.split(jax.vmap(self.wqkv)(x), 3, axis=-1)
        # (T, D) -> (n_heads, T, head_dim)
        split = lambda a: a.reshape(seq_len, self.n_heads, self.head_dim).transpose(
            1, 0, 2
        )
        q, k, v = split(q), split(k), split(v)
        scores = q @ k.transpose(0, 2, 1) / jnp.sqrt(self.head_dim)  # (H, T, T)
        mask = jnp.tril(jnp.ones((seq_len, seq_len)))
        scores = jnp.where(mask[None] == 0, -jnp.inf, scores)
        attn = jax.nn.softmax(scores, axis=-1)
        out = (attn @ v).transpose(1, 0, 2).reshape(seq_len, dim)  # back to (T, D)
        return jax.vmap(self.wo)(out)

Feed-forward (per-position MLP)

Linear(D -> 4D) -> GELU -> Linear(4D -> D), applied identically and independently to every position.

class FeedForward(eqx.Module):
    """Per-position MLP: ``Linear(dim, hidden)``, GELU, ``Linear(hidden, dim)``."""

    fc1: eqx.nn.Linear
    fc2: eqx.nn.Linear

    def __init__(self, dim, hidden, key):
        k1, k2 = random.split(key)
        self.fc1 = eqx.nn.Linear(dim, hidden, key=k1)
        self.fc2 = eqx.nn.Linear(hidden, dim, key=k2)

    def __call__(self, x):
        """Apply the MLP to every position of ``x`` of shape ``(T, D)``."""
        return jax.vmap(self.fc2)(jax.nn.gelu(jax.vmap(self.fc1)(x)))

Transformer block and full GPT

Pre-norm + residual around each sub-layer, then stack L blocks between the embeddings and the unembedding head.

class Block(eqx.Module):
    """One Transformer block: pre-norm attention and feed-forward sub-layers."""

    attn: MultiHeadAttention
    ff: FeedForward
    n1: eqx.nn.LayerNorm
    n2: eqx.nn.LayerNorm

    def __init__(self, dim, n_heads, hidden, key):
        k1, k2 = random.split(key)
        self.attn = MultiHeadAttention(dim, n_heads, k1)
        self.ff = FeedForward(dim, hidden, k2)
        self.n1 = eqx.nn.LayerNorm(dim)
        self.n2 = eqx.nn.LayerNorm(dim)

    def __call__(self, x):
        """Run both sub-layers on ``x`` of shape ``(T, D)`` through their residuals."""
        x = x + self.attn(jax.vmap(self.n1)(x))
        x = x + self.ff(jax.vmap(self.n2)(x))
        return x


class GPT(eqx.Module):
    """Decoder-only Transformer mapping token ids to logits over the vocabulary."""

    tok_emb: eqx.nn.Embedding
    pos_emb: eqx.nn.Embedding
    blocks: list
    norm_f: eqx.nn.LayerNorm
    head: eqx.nn.Linear
    context_length: int = eqx.field(static=True)

    def __init__(self, vocab_size, dim, n_heads, hidden, n_layers, context_length, key):
        keys = random.split(key, n_layers + 3)
        self.tok_emb = eqx.nn.Embedding(vocab_size, dim, key=keys[0])
        self.pos_emb = eqx.nn.Embedding(context_length, dim, key=keys[1])
        self.blocks = [
            Block(dim, n_heads, hidden, keys[2 + i]) for i in range(n_layers)
        ]
        self.norm_f = eqx.nn.LayerNorm(dim)
        self.head = eqx.nn.Linear(dim, vocab_size, use_bias=False, key=keys[-1])
        self.context_length = context_length

    def __call__(self, idx):
        """Map token ids ``(T,)`` to logits of shape ``(T, vocab_size)``."""
        seq_len = idx.shape[0]
        x = jax.vmap(self.tok_emb)(idx) + jax.vmap(self.pos_emb)(jnp.arange(seq_len))
        for block in self.blocks:
            x = block(x)
        return jax.vmap(self.head)(jax.vmap(self.norm_f)(x))

DIM, HEADS, HIDDEN, LAYERS, BATCH = 64, 4, 128, 2, 16

reference = GPT(vocab_size, DIM, HEADS, HIDDEN, LAYERS, T, key=random.key(0))

n_params = sum(
    x.size for x in jax.tree_util.tree_leaves(eqx.filter(reference, eqx.is_array))
)
print(f"number of parameters: {n_params:,}")

# sanity check: forward pass shape
logits = reference(data[:T])
print("logits shape:", logits.shape, "(T, vocab_size)")

number of parameters: 72,064

logits shape: (32, 27) (T, vocab_size)

Training with backpropagation

Next-token prediction with cross-entropy loss and Adam. eqx.filter_value_and_grad differentiates only the array leaves (the weights), leaving the static config untouched. This loss.backward()-style global gradient (one backward sweep threading information from the loss through every layer) is precisely what the PyHGF port replaces with a local, layer-wise prediction-error rule.

def ce_loss(model, x, y):
    """Mean cross-entropy of the model on a batch ``x`` with targets ``y``."""
    logits = jax.vmap(model)(x)  # (B, T, vocab_size)
    return optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()


optimiser = optax.adam(3e-3)
opt_state = optimiser.init(eqx.filter(reference, eqx.is_array))


@eqx.filter_jit
def backprop_step(model, opt_state, x, y):
    """Take one Adam step and return the updated model, optimiser state and loss."""
    loss, grads = eqx.filter_value_and_grad(ce_loss)(model, x, y)
    updates, opt_state = optimiser.update(grads, opt_state)
    return eqx.apply_updates(model, updates), opt_state, loss


N_STEPS = 400

key = random.key(1)
reference_losses = []
t0 = time.perf_counter()
for it in range(N_STEPS):
    key, bk = random.split(key)
    x, y = get_batch(bk, BATCH)
    reference, opt_state, loss = backprop_step(reference, opt_state, x, y)
    reference_losses.append(float(loss))
    if it % 100 == 0 or it == N_STEPS - 1:
        print(f"step {it:4d}  |  loss {float(loss):.4f}")
backprop_seconds = time.perf_counter() - t0
print(f"backpropagation: {N_STEPS} steps in {backprop_seconds:.0f} s")

2026-09-01 21:36:33.852495: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:33.865198: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:33.886694: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:33.898685: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:33.922052: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:33.935205: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:33.953063: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:33.965404: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:33.983070: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:33.994871: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.017465: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.030730: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.049207: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

step    0  |  loss 3.4361

2026-09-01 21:36:34.062543: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.080983: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.093962: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.112344: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.124248: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.143252: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.155412: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.173596: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.185618: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.205528: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.217419: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.240431: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.252544: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:34.271314: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.284601: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.302728: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.314741: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.333682: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.345398: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.363432: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.375466: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.394111: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.406325: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.424951: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.436760: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.457953: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.469888: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:34.488432: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.504471: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.522423: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.535539: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.553478: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.565474: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.584686: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.596788: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.615123: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.626878: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.645050: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.657291: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.680700: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:34.693032: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.712300: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.724716: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.744648: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.757845: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.776355: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.789686: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.807665: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.819805: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.839949: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.853478: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.872158: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.884944: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:34.904036: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.920252: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.938815: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.954096: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.973871: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:34.986105: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.004024: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.016289: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.035717: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.048029: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.066834: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.078789: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.101837: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:35.114295: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.133819: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.146248: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.164878: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.176596: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.194507: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.206760: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.225789: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.237486: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.255524: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.267677: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.285784: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.297754: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:35.316542: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.329628: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.348690: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.362505: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.380861: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.393009: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.411347: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.423581: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.441681: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.453627: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.471551: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.483959: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.501522: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.515953: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:35.539037: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.552954: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.572036: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.583663: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.602296: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.614250: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.632204: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.644081: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.661897: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.673885: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.692305: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.705443: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.724641: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.736210: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:35.755848: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.768133: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.786316: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.798492: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.816472: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.828798: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.848323: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.860226: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.878379: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.889809: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.907610: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.919974: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.939458: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.951700: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:35.971258: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:35.984810: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.004154: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.016508: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.035004: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.046884: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.067228: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.079159: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.097352: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.109282: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.129791: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.142175: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.161808: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:36.174070: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.193646: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.205672: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.224618: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.236620: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.255759: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.267507: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.286107: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.297846: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.316455: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.328594: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.348200: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.361141: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:36.380851: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.392585: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.411801: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.423774: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.443617: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.457157: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.476140: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.489671: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.509685: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.521784: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.541886: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.554151: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.574254: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:36.587613: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.606061: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.619320: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.639261: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.651512: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.669903: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.681730: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.700985: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.714278: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.732591: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.744581: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.762505: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.774684: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:36.792687: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.805228: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.825272: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.837145: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.855874: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.868009: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.886903: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.898603: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.917644: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.929741: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.949602: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.962223: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:36.986142: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:36.998800: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.018018: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.029817: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.049246: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.061512: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.081244: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.093383: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.112174: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.125511: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.144871: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.158009: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.176418: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.188560: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

step  100  |  loss 0.3172

2026-09-01 21:36:37.208295: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.221587: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.241457: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.254096: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.272405: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.284534: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.304489: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.316233: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.336642: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.348421: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.370982: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.382815: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.401831: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:37.414260: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.432777: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.446254: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.465445: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.477813: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.498556: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.514674: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.534743: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.546980: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.565560: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.577629: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.595857: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.607662: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:37.626523: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.638490: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.657851: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.669757: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.688627: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.700448: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.720241: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.733495: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.753006: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.766365: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.786065: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.797760: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.816119: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:37.830234: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.849514: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.861456: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.880123: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.892060: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.911963: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.923717: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.942288: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.954392: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.973176: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:37.986451: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.005403: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.017189: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:38.035442: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.047410: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.065646: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.077836: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.096689: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.108356: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.127170: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.139407: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.158030: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.170427: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.190430: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.202419: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.220577: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.233380: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:38.252789: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.264901: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.283299: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.295187: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.312982: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.326472: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.344876: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.356678: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.375121: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.387754: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.406678: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.422539: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.441853: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:38.454494: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.474064: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.485807: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.504887: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.517405: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.535711: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.551862: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.570891: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.582568: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.601450: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.614266: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.633762: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.647328: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:38.666056: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.679371: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.698842: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.713066: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.730992: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.743076: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.764321: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.776602: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.794578: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.806502: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.824315: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.836315: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.854680: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:38.869051: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.887989: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.900286: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.918616: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.930600: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.949427: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.962481: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.980479: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:38.992007: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.012101: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.023781: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.043687: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.055336: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:39.074995: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.086959: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.105079: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.117092: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.135074: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.147835: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.165824: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.177837: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.196054: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.209414: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.227896: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.240148: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.259619: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.271474: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:39.290274: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.302490: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.321484: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.333601: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.352251: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.364466: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.383011: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.397059: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.416561: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.428483: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.447638: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.460810: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.479901: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:39.492729: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.511731: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.523923: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.544029: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.555834: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.574228: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.586466: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.606020: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.619882: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.641148: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.653130: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.671860: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.683864: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:39.705369: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.717992: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.741073: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.752996: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.771959: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.783860: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.802314: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.814123: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.833060: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.845101: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.864780: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.876506: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.895150: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:39.907747: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.926723: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.940563: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.959428: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.971719: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:39.990744: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.003419: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.021771: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.033792: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.053967: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.066114: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.086551: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.101039: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:40.121658: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.133981: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.152295: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.164656: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.183508: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.195714: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.214525: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.226389: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.245379: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.257435: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.276459: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.288295: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.307137: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.318884: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

step  200  |  loss 0.1258

2026-09-01 21:36:40.338239: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.351169: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.370496: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.382578: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.401904: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.414679: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.433485: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.445638: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.465281: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.478117: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.497357: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.509759: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.528814: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:40.541604: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.560336: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.572673: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.592966: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.604870: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.623646: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.635608: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.654152: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.666214: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.684983: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.696844: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.715364: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.727579: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:40.746494: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.758406: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.777142: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.788804: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.809126: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.822532: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.842833: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.855078: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.874043: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.886390: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.904886: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.916794: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.935283: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.948883: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:40.968330: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.980300: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:40.999126: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.015267: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.033730: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.045698: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.064211: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.076354: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.096662: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.108569: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.126983: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.141491: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.160333: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:41.173164: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.195645: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.207460: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.225745: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.242057: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.261449: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.273163: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.291047: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.302776: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.320904: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.332777: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.350834: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.362978: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:41.381974: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.394127: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.412728: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.424575: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.442767: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.454864: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.477437: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.489741: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.509080: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.520995: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.539265: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.551649: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.570607: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.582344: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:41.601831: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.618317: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.636854: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.648773: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.667324: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.679459: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.698280: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.710076: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.728543: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.740444: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.759230: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.771241: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.789870: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:41.803502: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.823540: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.835982: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.855509: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.867373: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.886023: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.899612: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.918253: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.930164: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.949758: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.961862: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.979981: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:41.992266: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:42.012375: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.024265: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.042394: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.055683: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.074127: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.085906: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.104106: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.116057: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.134410: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.146006: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.165544: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.178501: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.196974: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.208724: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:42.227715: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.241178: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.261153: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.272847: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.290506: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.302786: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.321060: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.334817: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.352774: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.365223: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.383477: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.395319: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.415006: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.427125: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:42.447140: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.459254: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.482346: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.494190: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.512479: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.524454: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.544117: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.557456: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.576711: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.588518: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.608362: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.621542: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.640710: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:42.652812: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.671669: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.683746: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.702826: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.715002: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.734094: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.745980: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.764247: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.777921: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.799485: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.811664: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.833792: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.845724: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:42.865753: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.878938: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.897026: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.908922: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.928036: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.942009: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.960671: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.973922: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:42.992868: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.004977: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.023652: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.036274: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.055212: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:43.068445: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.090305: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.102750: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.121314: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.134467: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.152999: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.166246: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.184808: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.196528: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.215854: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.229037: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.249118: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.261025: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:43.279379: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.292721: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.311104: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.323353: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.342681: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.355006: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.373841: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.385822: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.407087: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.420404: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.439169: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.451272: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.470492: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

step  300  |  loss 0.0915

2026-09-01 21:36:43.482867: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.501165: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.514798: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.533375: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.545102: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.563231: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.576988: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.595119: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.608372: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.628011: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.640002: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.659425: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.671216: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:43.694065: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.707500: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.726767: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.738864: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.758719: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.772081: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.791111: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.802884: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.821404: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.833171: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.852778: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.864808: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.887136: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:43.903184: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.923427: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.935545: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.958974: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.970748: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:43.989863: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.001546: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.019730: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.033392: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.053136: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.066159: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.084660: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.097896: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:44.116426: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.128616: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.147980: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.160187: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.180490: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.196708: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.215425: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.228544: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.246743: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.258625: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.277647: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.289531: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.307993: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:44.320458: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.339894: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.352172: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.370822: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.382655: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.402295: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.414191: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.432390: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.445447: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.465526: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.478860: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.497498: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.513624: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:44.533093: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.546542: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.565605: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.578912: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.597207: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.610464: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.629491: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.641329: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.659961: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.672140: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.691162: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.703505: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.723151: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:44.735581: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.754196: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.766181: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.784702: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.796475: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.815414: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.827648: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.848046: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.861909: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.881467: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.893579: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.913019: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.925203: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:44.944870: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.957295: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.976041: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:44.988200: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.008150: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.021625: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.041137: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.053353: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.071750: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.083878: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.103530: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.116844: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.135323: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.147184: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:45.166623: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.178797: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.197362: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.209124: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.228028: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.241469: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.259834: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.272577: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.292044: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.303818: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.326420: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.338434: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.358143: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:45.370520: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.388910: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.400634: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.420598: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.433738: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.453580: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.466241: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.484695: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.496538: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.515649: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.527759: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.547554: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.559275: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:45.578203: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.589979: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.608136: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.622015: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.641820: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.653633: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.673267: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.685546: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.704143: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.716411: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.736031: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.749303: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.767266: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:45.779259: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.798382: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.810237: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.829372: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.841183: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.860081: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.872240: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.891116: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.903392: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.921683: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.933892: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.953159: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.965102: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:45.984906: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:45.997103: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.015491: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.027485: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.046178: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.058126: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.077228: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.089449: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.108421: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.120523: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.139000: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.150986: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.169628: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

2026-09-01 21:36:46.185997: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.205111: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.218270: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.236376: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.248628: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.267836: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.279842: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.299059: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.311014: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.330385: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.342581: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.362251: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.374064: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

step  399  |  loss 0.0804
backpropagation: 400 steps in 19 s

2026-09-01 21:36:46.393540: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.405988: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.424503: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.437012: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.454730: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.466735: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.485567: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed
2026-09-01 21:36:46.497555: W external/xla/xla/service/cpu/onednn_matmul.cc:333] [Perf]: MatMul reference implementation being executed

The same Transformer, learning through PyHGF

Now every weight in the model learns through PyHGF: the four attention tables of each block, the feed-forwards, both embedding tables, and the output head. Only the weightless calculations stay frozen: normalisations, the attention mixing, and the shortcut junctions, which translate errors with hand-derived formulas.

The pinned-confidence configuration

Every learning part keeps the same pinned-confidence configuration: volatility levels frozen and high prior confidence (precision, a belief’s inverse variance) on the hidden and input layers, so their beliefs barely move during the update sweep (a silent interior; see the theory page), with unit precision on the observed output layer.

The learning rule

The weights learn here in the exact backprop-parity mode (learning_kind="precision_weighted", the adapter’s default): the hidden layer’s belief shift is its routed error divided by its posterior confidence, and the confidence-weighted update multiplies that same confidence back in. The two cancel exactly, so each local update coincides with the backpropagated gradient (the correspondence is derived in the theory page). This makes the comparison with the baseline as direct as possible: the two learners apply the same weight updates and differ only in how they reach them, local message passing against a global backward sweep.

The categorical head

The head is a categorical belief layer: its nodes jointly represent one softmax choice over the characters, so the model optimises exactly the same cross-entropy objective as the baseline. Training clamps the truth “the next character is this one”, a one-hot pattern, directly onto those beliefs, and the classification error arises from PyHGF’s own prediction-error rule. The head’s own update is the plain cross-entropy gradient (the one_hot − softmax error times the activity), exactly the backprop head’s. The head’s forward output is already a probability distribution over the characters.

PARITY = dict(
    volatility_parent=False,
    precision=1e4,
    expected_precision=1e4,
)
PARITY_LEAF = dict(volatility_parent=False)

scratch = GPT(vocab_size, DIM, HEADS, HIDDEN, LAYERS, T, key=random.key(0))


def pyhgf_part(net):
    """Wrap a deep network as a pipeline part that learns through PyHGF."""
    # "precision_weighted" is the exact backprop-parity mode: each local update
    # coincides with the backpropagated gradient.
    return DeepNetworkAdapter(
        net, optimiser=optax.adam(3e-3), learning_kind="precision_weighted"
    )


pc_gpt = hybrid_from_gpt(
    scratch,
    attention_parts=[
        {
            name: pyhgf_part(
                from_linear(
                    getattr(b.attn, name), leaf_kwargs=PARITY_LEAF, layer_kwargs=PARITY
                )
            )
            for name in ("wqkv", "wo")
        }
        for b in scratch.blocks
    ],
    ff_parts=[
        pyhgf_part(
            from_feedforward(
                b.ff.fc1, b.ff.fc2, leaf_kwargs=PARITY_LEAF, layer_kwargs=PARITY
            )
        )
        for b in scratch.blocks
    ],
    head_part=pyhgf_part(
        from_linear(
            scratch.head,
            leaf_kwargs=dict(kind="categorical", **PARITY_LEAF),
            layer_kwargs=PARITY,
        )
    ),
    token_part=pyhgf_part(
        from_embedding(scratch.tok_emb, leaf_kwargs=PARITY_LEAF, layer_kwargs=PARITY)
    ),
    position_part=pyhgf_part(
        from_embedding(scratch.pos_emb, leaf_kwargs=PARITY_LEAF, layer_kwargs=PARITY)
    ),
)

pc_pipeline = FusedPipeline(
    pc_gpt,
    # The cross-entropy error at the head, formed inside the compiled step:
    # feeding probs - one_hot through the correct-and-clamp entry clamps
    # exactly the one-hot truth on the beliefs.
    error_fn=lambda probs, y: probs - jax.nn.one_hot(y, vocab_size),
)

# The timed loop does the same work as the backpropagation loop (draw a
# batch, take one step) and only stores each step's output; every metric
# is computed after the fact, below.
all_probs, all_targets = [], []
key = random.key(4)
t0 = time.perf_counter()
for it in range(N_STEPS):
    key, bk = random.split(key)
    x, y = get_batch(bk, BATCH)
    probs, _ = pc_pipeline.step(x, y)  # (B, T, vocab): one softmax per position
    all_probs.append(probs)
    all_targets.append(y)
probs.block_until_ready()  # flush the queued steps before stopping the clock
pyhgf_seconds = time.perf_counter() - t0
print(f"PyHGF: {N_STEPS} steps in {pyhgf_seconds:.0f} s")

# Metrics, after the fact. The head's output is a probability distribution,
# so its cross-entropy is directly comparable with the reference's.
all_probs, all_targets = jnp.stack(all_probs), jnp.stack(all_targets)
p_target = jnp.take_along_axis(all_probs, all_targets[..., None], axis=-1)
pc_losses = -jnp.log(jnp.clip(p_target, 1e-9, None)).mean(axis=(1, 2, 3))
accuracies = (all_probs.argmax(-1) == all_targets).mean(axis=(1, 2))
for it in (0, 100, 200, 300, N_STEPS - 1):
    print(f"step {it:4d}  |  next-character accuracy {float(accuracies[it]):.2f}")

PyHGF: 400 steps in 14 s

step    0  |  next-character accuracy 0.06
step  100  |  next-character accuracy 0.91
step  200  |  next-character accuracy 0.96
step  300  |  next-character accuracy 0.98
step  399  |  next-character accuracy 0.97

fig, ax = plt.subplots(figsize=(7, 4))
ax.plot(accuracies, color="#4c72b0", linewidth=1.5)
ax.set(
    xlabel="training step",
    ylabel="next-character accuracy",
    title="Every weight learning through PyHGF, no backpropagation anywhere",
)
ax.grid(linestyle="--")
sns.despine()

Accuracy and execution time

Both models are trained for the same number of steps on batches drawn in the same way, and both run as one compiled program per step, so their training curves and wall-clock times can be put side by side. With the categorical head, both learners optimise the same softmax cross-entropy, so the curves are directly comparable — and in the backprop-parity mode used here each local update coincides with the backpropagated gradient.

fig, (ax1, ax2) = plt.subplots(
    1, 2, figsize=(10, 4), gridspec_kw={"width_ratios": [2.2, 1]}
)

ax1.plot(
    reference_losses, label="Backpropagation (Equinox)", color="grey", linewidth=1.5
)
ax1.plot(pc_losses, label="Predictive coding (PyHGF)", color="#4c72b0", linewidth=1.5)
ax1.set(
    xlabel="Training step",
    ylabel="Training cross-entropy",
    title="Training loss per step",
)
ax1.grid(linestyle="--")
ax1.legend()

bars = ax2.bar(
    ["Backpropagation \n (Equinox)", "Predictive coding \n (PyHGF)"],
    [backprop_seconds, pyhgf_seconds],
    color=["grey", "#4c72b0"],
)

ax2.set(ylabel=f"Training time, {N_STEPS} steps (s)", title="Wall-clock cost")
ax2.grid(axis="y", linestyle="--")

sns.despine()

ms_backprop = backprop_seconds / N_STEPS * 1e3
ms_pyhgf = pyhgf_seconds / N_STEPS * 1e3
print(
    f"time per step: backprop {ms_backprop:.1f} ms, "
    f"PyHGF {ms_pyhgf:.1f} ms ({ms_pyhgf / ms_backprop:.1f}x)"
)

time per step: backprop 46.4 ms, PyHGF 34.3 ms (0.7x)

In terms of learning per step, the two are close: both collapse the loss within the first few hundred steps and flatten near the same floor. In terms of wall-clock, the predictive-coding learner stays within a small factor of backpropagation.

Generating text

Autoregressive generation is an inference-time loop around the feedforward model: feed the current context (cropped to the context length), take the prediction at the last position, sample the next token, append, and repeat. This outer loop is not recurrence in the architecture; the model itself remains feedforward. Both learners generate the same way; the only difference is where the next-character distribution comes from. The backprop model outputs logits:

def generate_backprop(model, idx, n_new, key):
    """Sample ``n_new`` characters from the backpropagation-trained model."""
    for _ in range(n_new):
        logits = model(idx[-model.context_length :])[-1]  # last-position logits
        key, sub = random.split(key)
        next_id = random.categorical(sub, logits).astype(jnp.int32)
        idx = jnp.concatenate([idx, next_id[None]])
    return idx


out = generate_backprop(reference, encode("predictive"), n_new=120, key=random.key(2))
print(decode(out))

predictive coding networks learn by minimizing prediction errors. each layer predicts the activity of the next, and updates its we

The PyHGF model’s forward pass already ends in a probability distribution over the characters (its categorical head), which is sampled from directly:

def generate_pyhgf(pipeline, idx, n_new, key):
    """Sample ``n_new`` characters from the categorical head of the pipeline."""
    for _ in range(n_new):
        window = idx[-T:]
        probs = pipeline.predict(window[None])[0]  # forward pass only
        p = probs[window.shape[0] - 1]
        key, sub = random.split(key)
        next_id = random.categorical(sub, jnp.log(p)).astype(jnp.int32)
        idx = jnp.concatenate([idx, next_id[None]])
    return idx


out = generate_pyhgf(pc_pipeline, encode("predictive"), n_new=120, key=random.key(5))
print(decode(out))

predictive coding networks learn by minimizing prediction errors. each layer predicts the activity of the next, and updates its we

System configuration

%load_ext watermark
%watermark -n -u -v -iv -w -p pyhgf,jax,jaxlib

Last updated: Tue, 01 Sep 2026

Python implementation: CPython
Python version       : 3.12.14
IPython version      : 9.16.1

pyhgf : 0.3.1
jax   : 0.6.2
jaxlib: 0.6.2

equinox   : 0.13.8
jax       : 0.6.2
matplotlib: 3.11.1
optax     : 0.2.8
pyhgf     : 0.3.1
seaborn   : 0.13.2

Watermark: 2.6.0