anima_closeout v15: the val bed was reading the eval cache wrong in TWO ways. (1) text embeddings live in the AR-bucket dir (ar_frames_{ar}x{f}, dataset.py:485), not the size-bucket dir (cache_{ar}x{w}x{h}x{f}, :247) — hence 'found []'. (2) far worse: SizeBucketDataset SHUFFLES its metadata (:241-242) before caching latents while TE is cached unshuffled one level up, so pairing by row order marries each latent to someone else's caption and still yields plausible numbers. The bed now joins on image_spec — the key the fork itself joins on (:312) — and iterates in sorted spec order for run-to-run reproducibility. PROVEN on a synthetic cache written by the fork's OWN ParquetCache with the real shuffle: 24/24 rows correctly paired, while the old order-based read mismatched 23/24.

colab/anima_closeout.ipynb

@@ -954,19 +954,22 @@
     "DEV = \"cuda\"\n",
     "\n",
     "# -- read the fork's eval cache -------------------------------------------\n",
-    "# Layout (utils/dataset.py + parquet_cache.py): {root}/cache_{bucket}/ holds\n",
-    "# separate parquet caches \u2014 'latents' (models/base.py:247) and\n",
-    "# 'text_embeddings_{i}' (cosmos_predict2.py:398: prompt_embeds/attn_mask/\n",
-    "# t5_input_ids/t5_attn_mask; the live anima runs write i=1, so the TE dir\n",
-    "# is DISCOVERED, not hardcoded; 'uncond_text_embeddings_*' is excluded).\n",
-    "# Both map over the SAME shuffled metadata dataset and this campaign has\n",
-    "# exactly one caption per image, so rows pair by order; the per-bucket\n",
-    "# count assert guards that assumption.\n",
+    "# LAYOUT (utils/dataset.py) and the trap inside it:\n",
+    "#   {root}/cache_{ar}x{w}x{h}x{f}/latents/         SizeBucketDataset (:247)\n",
+    "#   {root}/ar_frames_{ar}x{f}/text_embeddings_N/   ARBucketDataset   (:485)\n",
+    "# Latents are cached from metadata that SizeBucketDataset SHUFFLES first\n",
+    "# (:241-242); text embeddings are cached one level up, at the AR bucket,\n",
+    "# unshuffled. So latent row i and TE row i are DIFFERENT images \u2014 pairing\n",
+    "# by row order would silently hand every latent someone else's caption and\n",
+    "# still produce plausible numbers. Both caches carry image_spec\n",
+    "# (parquet_cache.image_specs()), which is the key the fork itself joins on\n",
+    "# (dataset.py:312), so we join on image_spec and iterate in sorted spec\n",
+    "# order, which also makes the bed reproducible run-to-run.\n",
     "cache_root = CFG[\"eval_cache_dir\"]\n",
-    "buckets = sorted(d for d in os.listdir(cache_root)\n",
-    "                 if d.startswith(\"cache_\")\n",
-    "                 and os.path.isdir(os.path.join(cache_root, d)))\n",
-    "assert buckets, f\"no cache_* bucket dirs under {cache_root}\"\n",
+    "size_dirs = sorted(d for d in os.listdir(cache_root)\n",
+    "                   if d.startswith(\"cache_\")\n",
+    "                   and os.path.isdir(os.path.join(cache_root, d)))\n",
+    "assert size_dirs, f\"no cache_* size-bucket dirs under {cache_root}\"\n",
     "\n",
     "def read_cache(d):\n",
     "    shards = sorted(p for p in os.listdir(d)\n",
@@ -981,29 +984,56 @@
     "    return torch.frombuffer(bytearray(raw),\n",
     "                            dtype=getattr(torch, dt)).reshape(shape)\n",
     "\n",
+    "def by_spec(tables, what):\n",
+    "    \"\"\"{image_spec tuple: (table, row)} \u2014 tuples are JSON-encoded by the\n",
+    "    cache writer (parquet_cache.py:_encode_row), so accept either form.\"\"\"\n",
+    "    out = {}\n",
+    "    for t in tables:\n",
+    "        names = t.schema.names\n",
+    "        col = (\"image_spec__json\" if \"image_spec__json\" in names else\n",
+    "               (\"image_spec\" if \"image_spec\" in names else None))\n",
+    "        assert col, f\"{what} cache carries no image_spec column: {names[:8]}\"\n",
+    "        for j, v in enumerate(t.column(col).to_pylist()):\n",
+    "            s = tuple(json.loads(v)) if isinstance(v, str) else tuple(v)\n",
+    "            out.setdefault(s, (t, j))\n",
+    "    return out\n",
+    "\n",
     "LAT = [\"latents\"]\n",
     "TE = [\"prompt_embeds\", \"attn_mask\", \"t5_input_ids\", \"t5_attn_mask\"]\n",
     "rows = []\n",
-    "for b in buckets:\n",
-    "    lt = read_cache(os.path.join(cache_root, b, \"latents\"))\n",
-    "    te_dirs = sorted(d for d in os.listdir(os.path.join(cache_root, b))\n",
-    "                     if d.startswith(\"text_embeddings_\"))\n",
-    "    assert len(te_dirs) == 1, \\\n",
-    "        f\"{b}: expected exactly one text_embeddings_* cache, found {te_dirs}\"\n",
-    "    tt = read_cache(os.path.join(cache_root, b, te_dirs[0]))\n",
+    "for sd in size_dirs:\n",
+    "    parts = sd[len(\"cache_\"):].split(\"x\")        # {ar}x{w}x{h}x{frames}\n",
+    "    ar_dir = os.path.join(cache_root, f\"ar_frames_{parts[0]}x{parts[-1]}\")\n",
+    "    assert os.path.isdir(ar_dir), (\n",
+    "        f\"{sd}: expected AR-bucket dir {ar_dir} \u2014 text embeddings are \"\n",
+    "        \"cached per AR bucket, not per size bucket\")\n",
+    "    te_names = sorted(d for d in os.listdir(ar_dir)\n",
+    "                      if d.startswith(\"text_embeddings_\"))\n",
+    "    assert len(te_names) == 1, \\\n",
+    "        f\"{ar_dir}: expected one text_embeddings_* cache, found {te_names}\"\n",
+    "    lat = by_spec(read_cache(os.path.join(cache_root, sd, \"latents\")),\n",
+    "                  \"latents\")\n",
+    "    te = by_spec(read_cache(os.path.join(ar_dir, te_names[0])),\n",
+    "                 \"text embedding\")\n",
     "    for c in LAT:\n",
-    "        assert c + \"__shape\" in lt[0].schema.names, (b, lt[0].schema.names)\n",
+    "        _s = next(iter(lat.values()))[0].schema.names\n",
+    "        assert c + \"__shape\" in _s, (sd, _s[:8])\n",
     "    for c in TE:\n",
-    "        assert c + \"__shape\" in tt[0].schema.names, (b, tt[0].schema.names)\n",
-    "    lrows = [(t, j) for t in lt for j in range(t.num_rows)]\n",
-    "    trows = [(t, j) for t in tt for j in range(t.num_rows)]\n",
-    "    assert len(lrows) == len(trows), \\\n",
-    "        f\"{b}: latents {len(lrows)} != text {len(trows)} \u2014 pairing unsafe\"\n",
-    "    for (ltb, lj), (ttb, tj) in zip(lrows, trows):\n",
-    "        r = {c: decol(ltb, c, lj) for c in LAT}\n",
-    "        r.update({c: decol(ttb, c, tj) for c in TE})\n",
+    "        _s = next(iter(te.values()))[0].schema.names\n",
+    "        assert c + \"__shape\" in _s, (ar_dir, _s[:8])\n",
+    "    _missing = [s for s in lat if s not in te]\n",
+    "    assert not _missing, (\n",
+    "        f\"{sd}: {len(_missing)} latents have no text embedding for their \"\n",
+    "        f\"image_spec (e.g. {_missing[:1]}) \u2014 caches are out of sync\")\n",
+    "    for s in sorted(lat):\n",
+    "        lt, lj = lat[s]\n",
+    "        tt, tj = te[s]\n",
+    "        r = {c: decol(lt, c, lj) for c in LAT}\n",
+    "        r.update({c: decol(tt, c, tj) for c in TE})\n",
     "        rows.append(r)\n",
-    "print(f\"eval cache rows: {len(rows)} from {len(buckets)} bucket dir(s)\",\n",
+    "    print(f\"  {sd}: {len(lat)} rows joined on image_spec \"\n",
+    "          f\"(AR pool {len(te)})\", flush=True)\n",
+    "print(f\"eval cache rows: {len(rows)} from {len(size_dirs)} size bucket(s)\",\n",
     "      flush=True)\n",
     "rows = rows[:int(os.environ.get(\"VB_MAX_ROWS\", \"512\"))]\n",
     "assert rows\n",