Upload completed Trace-iMF training run

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+diagnostics/dashboard.png filter=lfs diff=lfs merge=lfs -text
+diagnostics/multitrace.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 TowerClimber
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,44 @@
 ---
+library_name: pytorch
 license: mit
+tags: [meanflow, multi-trace, image-generation, class-conditional, cifar10]
 ---
+
+# Multi-Trace U-Net — 14,992,182 parameters
+
+U-Net adapted from the user-provided Untitled0.ipynb and widened to approximately
+22M parameters. Trained for 50,000 optimizer updates on CIFAR-10.
+Training samples K uniformly from 1..5 and a grid interval j/K;
+the same u head fits the interval and diagonal, with a detached directional JVP.
+Default inference weights are **ema**. Model input/output pixels
+are 32x32 RGB; sampling maps the final [-1,1] image into [0,1].
+
+This is a custom PyTorch model. Download with huggingface_hub.snapshot_download,
+install the included requirements.txt, then run:
+
+```bash
+python sample.py --model-dir . --nfe 1 --output samples_1nfe.png
+python sample.py --model-dir . --nfe 5 --output samples_5nfe.png
+```
+
+config.json describes architecture/objective; training_config.json records the
+effective batch, accumulation, LR schedule, clipping and EMA decay.
+Full optimizer/EMA/RNG resume checkpoints remain on the Modal Volume and, when
+the Modal notebook backup option is enabled, are also stored under
+training-checkpoints/<run-name>/ in this Hub repository.
+best_validation.safetensors contains the selected EMA validation snapshot;
+its step and fixed held-out losses are recorded in best_validation.json.
+Diagnostics include gradient RMS, actual optimizer updates, losses by K,
+clipping frequency, samples, throughput, ETA and TensorBoard events.
+Plateau flags are heuristics, not proof of convergence or image quality.
+
+When present, fid.json reports pytorch-fid pool3 FID against 50,000 unaugmented
+CIFAR-10 training images, with the generated sample count and each NFE explicitly
+recorded. This protocol differs from the reference notebook's ImageNet-feature
+FID; the values must not be compared directly. No quality claim is made here.
+The test split was used as fixed validation during training.
+
+Labels: airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck.
+The U-Net and interval schedule follow the user-provided notebook; repository
+training flow and DiT baseline are adapted from haidog-yaqub/MeanFlow (MIT).
+See third_party/ for attribution.

best_validation.json

@@ -0,0 +1,53 @@
+{
+  "step": 50000,
+  "weights_source": "ema",
+  "loss": 0.3709731566905975,
+  "edge_loss": 0.19136443853378296,
+  "diag_loss": 0.1796087145805359,
+  "jvp_rms": 2.2817168676853177,
+  "examples": 5120,
+  "loss_nfe_1": 0.39239816665649413,
+  "edge_loss_nfe_1": 0.21810964345932007,
+  "diag_loss_nfe_1": 0.1742885172367096,
+  "jvp_rms_nfe_1": 2.700194501876831,
+  "loss_nfe_2": 0.3679115891456604,
+  "edge_loss_nfe_2": 0.18770996630191802,
+  "diag_loss_nfe_2": 0.18020162284374236,
+  "jvp_rms_nfe_2": 2.536956226825714,
+  "loss_nfe_3": 0.35463637113571167,
+  "edge_loss_nfe_3": 0.1791027843952179,
+  "diag_loss_nfe_3": 0.17553358674049377,
+  "jvp_rms_nfe_3": 2.0643750607967375,
+  "loss_nfe_4": 0.38753209710121156,
+  "edge_loss_nfe_4": 0.19491184055805205,
+  "diag_loss_nfe_4": 0.19262024462223054,
+  "jvp_rms_nfe_4": 2.4871247947216033,
+  "loss_nfe_5": 0.3523875594139099,
+  "edge_loss_nfe_5": 0.17698795795440675,
+  "diag_loss_nfe_5": 0.17539960145950317,
+  "jvp_rms_nfe_5": 1.6199337542057037,
+  "raw_loss": 0.37272677540779114,
+  "raw_edge_loss": 0.19294691383838652,
+  "raw_diag_loss": 0.17977986514568328,
+  "raw_jvp_rms": 2.4338310968875883,
+  "raw_loss_nfe_1": 0.39853897094726565,
+  "raw_edge_loss_nfe_1": 0.2240687519311905,
+  "raw_diag_loss_nfe_1": 0.1744702249765396,
+  "raw_jvp_rms_nfe_1": 2.9251581192016602,
+  "raw_loss_nfe_2": 0.3692192077636719,
+  "raw_edge_loss_nfe_2": 0.18882804214954377,
+  "raw_diag_loss_nfe_2": 0.18039117157459258,
+  "raw_jvp_rms_nfe_2": 2.7281934142112734,
+  "raw_loss_nfe_3": 0.35528815984725953,
+  "raw_edge_loss_nfe_3": 0.1796249717473984,
+  "raw_diag_loss_nfe_3": 0.17566319108009337,
+  "raw_jvp_rms_nfe_3": 2.1835286557674407,
+  "raw_loss_nfe_4": 0.3877676367759705,
+  "raw_edge_loss_nfe_4": 0.19502896666526795,
+  "raw_diag_loss_nfe_4": 0.1927386701107025,
+  "raw_jvp_rms_nfe_4": 2.6403229236602783,
+  "raw_loss_nfe_5": 0.3528199017047882,
+  "raw_edge_loss_nfe_5": 0.1771838366985321,
+  "raw_diag_loss_nfe_5": 0.17563606798648834,
+  "raw_jvp_rms_nfe_5": 1.69195237159729
+}

best_validation.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3e34b7db78ae343f152e18c05309ffdcaf74cf8b6bcb6f9d8d4bf98e86f260bd
+size 59989568

config.json

@@ -0,0 +1,38 @@
+{
+  "architecture": "ConditionalMultiTraceUNet",
+  "model": {
+    "input_size": 32,
+    "in_channels": 3,
+    "num_classes": 10,
+    "base_channels": 69,
+    "channel_mult": [
+      1,
+      2,
+      2,
+      4
+    ],
+    "num_res_blocks": 1,
+    "cond_dim": 312,
+    "fourier_dim": 64,
+    "attn_resolutions": [
+      8
+    ]
+  },
+  "objective": "multi_trace_imf",
+  "multi_trace_imf": {
+    "lambda_diag": 1.0,
+    "max_nfe": 5,
+    "jvp_precision": "bf16",
+    "channels_last": true
+  },
+  "dataset": "cifar10",
+  "weights_source": "ema",
+  "max_nfe": 5,
+  "training_steps": 50000,
+  "trainable_parameters": 14992182,
+  "pixel_range": [
+    0,
+    1
+  ],
+  "sampling_steps": 1
+}

diagnostics/convergence.json

@@ -0,0 +1,15 @@
+{
+  "status": "plateau_candidate_late",
+  "required_evaluations": 5,
+  "warning": "Regression loss has an irreducible noise floor; these are diagnostic hints, not FID.",
+  "step": 50000,
+  "best_validation_step": 50000,
+  "best_validation_loss": 0.3709731566905975,
+  "latest_validation_loss": 0.3709731566905975,
+  "relative_improvement": 0.0015668281478368013,
+  "window_start_step": 46000,
+  "window_end_step": 50000,
+  "budget_phase": "late",
+  "validation_regressed_from_best": false,
+  "recent_clip_fraction": 0.015999999642372132
+}

fid.json

@@ -0,0 +1,36 @@
+{
+  "protocol": {
+    "implementation": "pytorch-fid==0.3.0",
+    "features": "FID Inception pool3 (2048)",
+    "input_range": [
+      0,
+      1
+    ],
+    "real_split": "CIFAR-10 train, unaugmented",
+    "real_count": 50000,
+    "real_data_sha256": "47af8b02cfcfb6cdf3dbbd1969ad7f1b70f4413f6f80f437b020ee4e5bc8304a",
+    "generated_count": 10000,
+    "seed": 2026,
+    "batch_size": 128,
+    "class_schedule": "balanced labels: sample index modulo 10",
+    "sampling_precision": "bf16",
+    "weights_source": "ema",
+    "weights_sha256": "7e739a7d2fab764f4ec8c02dac518d6c4cabab9cd5fc33e4aaf410064f93b540",
+    "training_step": 50000,
+    "nfes": [
+      1,
+      2,
+      3,
+      4,
+      5
+    ],
+    "reference_notebook_comparable": false
+  },
+  "fid": {
+    "1": 38.301275424967,
+    "2": 26.30447400280991,
+    "3": 25.51997010873481,
+    "4": 25.017732838651114,
+    "5": 24.852908128389515
+  }
+}

hub_checkpoints.jsonl

@@ -0,0 +1 @@
+{"repo_id": "fushinguyenex/IMF", "run_name": "multitrace_unet15m_h200_b512_v3", "checkpoint": "step_0050000.pt", "step": 50000, "revision": "7487b07fea07a88eec41fa931164d65e55c6fddd", "path_in_repo": "training-checkpoints/multitrace_unet15m_h200_b512_v3/step_0050000.pt", "url": "https://huggingface.co/fushinguyenex/IMF/blob/7487b07fea07a88eec41fa931164d65e55c6fddd/training-checkpoints/multitrace_unet15m_h200_b512_v3/step_0050000.pt", "status": "uploaded", "attempts": 1}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e06188fade7ae05fb718146517aa998b19fe4f907c921d70d0521146cac0b709
+size 59989600

models/__init__.py

@@ -0,0 +1,10 @@
+from .dit import TraceDiT
+from .unet import ConditionalMultiTraceUNet
+
+
+def build_backbone(cfg):
+    architecture = cfg.get("architecture", "TraceDiT")
+    classes = {"TraceDiT": TraceDiT, "ConditionalMultiTraceUNet": ConditionalMultiTraceUNet}
+    if architecture not in classes:
+        raise ValueError(f"Unknown architecture: {architecture}")
+    return classes[architecture](**cfg["model"])

models/dit.py

@@ -0,0 +1,352 @@
+# Adapted from haidog-yaqub/MeanFlow, MIT; see third_party/MeanFlow.LICENSE.
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import math
+
+
+class RMSNorm(nn.Module):
+    """RMSNorm compatible with PyTorch 2.2 and forward-mode differentiation."""
+    def __init__(self, dim, elementwise_affine=True):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim)) if elementwise_affine else None
+
+    def forward(self, x):
+        dtype = x.dtype
+        value = x.float()
+        value = value * torch.rsqrt(value.square().mean(-1, keepdim=True) + 1e-6)
+        if self.weight is not None:
+            value = value * self.weight
+        return value.to(dtype)
+
+
+class PatchEmbed(nn.Module):
+    """Non-overlapping convolutional patch projection, equivalent to source timm usage."""
+    def __init__(self, input_size, patch_size, in_channels, dim):
+        super().__init__()
+        self.input_size = input_size
+        self.patch_size = (patch_size, patch_size)
+        self.num_patches = (input_size // patch_size) ** 2
+        self.proj = nn.Conv2d(in_channels, dim, patch_size, stride=patch_size)
+
+    def forward(self, x):
+        if x.shape[-2:] != (self.input_size, self.input_size):
+            raise ValueError(f"Expected {self.input_size}x{self.input_size} input, got {x.shape[-2:]}")
+        return self.proj(x).flatten(2).transpose(1, 2)
+
+
+class Mlp(nn.Sequential):
+    def __init__(self, in_features, hidden_features, act_layer, drop=0):
+        super().__init__(nn.Linear(in_features, hidden_features), act_layer(),
+                         nn.Linear(hidden_features, in_features))
+
+
+def expand_ada_params(param, x):
+    """Expand global (B, D) adaLN params to (B, T, D); pass through token-level params."""
+    if param.ndim == 2:
+        return param.unsqueeze(1)
+    return param
+
+
+def modulate(x, scale, shift):
+    """
+    Apply adaLN modulation.
+    x: (B, T, D)
+    scale, shift: (B, D) for global adaLN or (B, T, D) for token-level adaLN
+    """
+    scale = expand_ada_params(scale, x)
+    shift = expand_ada_params(shift, x)
+    return x * (1 + scale) + shift
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, dim, nfreq=256, scale=1000.0):
+        super().__init__()
+        self.mlp = nn.Sequential(nn.Linear(nfreq, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.nfreq = nfreq
+        self.scale = scale
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        half_dim = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period)
+            * torch.arange(start=0, end=half_dim, dtype=torch.float32)
+            / half_dim
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat(
+                [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
+            )
+        return embedding
+
+    def forward(self, t):
+        t = t * self.scale
+        t_freq = self.timestep_embedding(t, self.nfreq)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+    def initialize_weights(self):
+        nn.init.normal_(self.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.mlp[2].weight, std=0.02)
+
+
+class LabelEmbedder(nn.Module):
+    def __init__(self, num_classes, dim):
+        super().__init__()
+        self.embedding = nn.Embedding(num_classes, dim)
+        self.num_classes = num_classes
+
+    def forward(self, labels):
+        embeddings = self.embedding(labels)
+        return embeddings
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=True, qk_norm=True):
+        super().__init__()
+        assert dim % num_heads == 0
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.q_norm = RMSNorm(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = RMSNorm(self.head_dim) if qk_norm else nn.Identity()
+        self.proj = nn.Linear(dim, dim)
+
+    def _native_attention(self, q, k, v):
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        return attn @ v
+
+    def _flash_attention(self, q, k, v):
+        return F.scaled_dot_product_attention(q, k, v, scale=self.scale)
+
+    def forward(self, x, use_flash_attention=False):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        if (
+            use_flash_attention
+            and q.is_cuda
+            and q.dtype in (torch.float16, torch.bfloat16)
+        ):
+            x = self._flash_attention(q, k, v)
+        else:
+            x = self._native_attention(q, k, v)
+
+        x = x.transpose(1, 2).reshape(B, N, C)
+        return self.proj(x)
+
+
+class DiTBlock(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4.0):
+        super().__init__()
+        self.norm1 = RMSNorm(dim, elementwise_affine=False)
+        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=True, qk_norm=True)
+        self.norm2 = RMSNorm(dim, elementwise_affine=False)
+        mlp_dim = int(dim * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.mlp = Mlp(
+            in_features=dim, hidden_features=mlp_dim, act_layer=approx_gelu, drop=0
+        )
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
+
+    def forward(self, x, c, use_flash_attention=False):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            self.adaLN_modulation(c).chunk(6, dim=-1)
+        )
+        x = x + expand_ada_params(gate_msa, x) * self.attn(
+            modulate(self.norm1(x), scale_msa, shift_msa),
+            use_flash_attention=use_flash_attention,
+        )
+        x = x + expand_ada_params(gate_mlp, x) * self.mlp(
+            modulate(self.norm2(x), scale_mlp, shift_mlp)
+        )
+        return x
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, dim, patch_size, out_dim):
+        super().__init__()
+        self.norm_final = RMSNorm(dim, elementwise_affine=False)
+        self.linear = nn.Linear(dim, patch_size * patch_size * out_dim)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim))
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+        x = modulate(self.norm_final(x), scale, shift)
+        x = self.linear(x)
+        return x
+
+
+class TraceDiT(nn.Module):
+    def __init__(
+        self,
+        input_size=32,
+        patch_size=2,
+        in_channels=3,
+        dim=384,
+        depth=8,
+        num_heads=6,
+        mlp_ratio=4.0,
+        num_classes=10,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+        self.num_classes = num_classes
+
+        if input_size <= 0 or patch_size <= 0 or input_size % patch_size:
+            raise ValueError("input_size must be positive and divisible by patch_size")
+        if dim <= 0 or dim % 4 or num_heads <= 0 or dim % num_heads:
+            raise ValueError("dim must be positive and divisible by 4 and num_heads")
+        if depth <= 0 or mlp_ratio <= 0:
+            raise ValueError("depth and mlp_ratio must be positive")
+        self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, dim)
+        self.t_embedder = TimestepEmbedder(dim)
+        self.r_embedder = TimestepEmbedder(dim)
+
+        self.use_cond = num_classes is not None
+        self.y_embedder = LabelEmbedder(num_classes, dim) if self.use_cond else None
+
+        num_patches = self.x_embedder.num_patches
+        self.register_buffer("pos_embed", torch.zeros(1, num_patches, dim))
+
+        self.blocks = nn.ModuleList([
+            DiTBlock(dim, num_heads, mlp_ratio) for _ in range(depth)
+        ])
+        self.final_layer = FinalLayer(dim, patch_size, self.out_channels)
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+        # Initialize (and freeze) pos_embed by sin-cos embedding:
+        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.x_embedder.num_patches ** 0.5))
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
+        w = self.x_embedder.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.x_embedder.proj.bias, 0)
+
+        # Initialize label embedding table:
+        if self.y_embedder is not None:
+            nn.init.normal_(self.y_embedder.embedding.weight, std=0.02)
+
+        # Initialize timestep embedding MLP (t, r):
+        for embedder in (self.t_embedder, self.r_embedder):
+            embedder.initialize_weights()
+
+        # Zero-out adaLN modulation layers in DiT blocks:
+        for block in self.blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        # Zero-out output layers:
+        for final_layer in (self.final_layer,):
+            nn.init.constant_(final_layer.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(final_layer.adaLN_modulation[-1].bias, 0)
+            nn.init.constant_(final_layer.linear.weight, 0)
+            nn.init.constant_(final_layer.linear.bias, 0)
+
+    def unpatchify(self, x):
+        """
+        x: (N, T, patch_size**2 * C)
+        imgs: (N, H, W, C)
+        """
+        c = self.out_channels
+        p = self.x_embedder.patch_size[0]
+        h = w = int(x.shape[1] ** 0.5)
+        assert h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p))
+        return imgs
+
+    def forward(self, x, t, r, y=None, use_flash_attention=False):
+        """Return u(z,r,t); API preserves source order (x,t,r).
+
+        x: (B,C,H,W); t,r: (B,); y: class indices when conditioned.
+        The diagonal velocity uses this same head with r=t.
+        """
+        x = self.x_embedder(x) + self.pos_embed
+        c = self.t_embedder(t) + self.r_embedder(r)
+        if self.use_cond:
+            if y is None:
+                raise ValueError("Class-conditioned model requires labels y")
+            c = c + self.y_embedder(y)
+        for block in self.blocks:
+            x = block(x, c, use_flash_attention=use_flash_attention)
+        return self.unpatchify(self.final_layer(x, c))
+
+
+# Positional embedding from:
+# https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb

models/unet.py

@@ -0,0 +1,214 @@
+"""Conditional U-Net adapted from the user-provided Untitled0.ipynb.
+
+Architecture lineage: base channels 44, (1,2,2,4), one residual block per level.
+Uses the same (z,t,r,y) API as TraceDiT so training and Hub sampling share a path.
+"""
+import math
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+def gn_groups(ch, max_groups=8):
+    for g in range(min(max_groups, ch), 0, -1):
+        if ch % g == 0:
+            return g
+    return 1
+
+
+class JVPGroupNorm(nn.GroupNorm):
+    def forward(self, x):
+        # PyTorch 2.2 forward AD can fail on channels-last GroupNorm strides.
+        # Normalize a contiguous view; convolutions still use channels-last.
+        return super().forward(x.contiguous())
+
+
+class ScalarEmbed(nn.Module):
+    """Sinusoidal scalar embedding with cached frequencies."""
+    def __init__(self, out_dim, fourier_dim=64):
+        super().__init__()
+        self.fourier_dim = fourier_dim
+        half = fourier_dim // 2
+        freqs = torch.exp(
+            -math.log(10_000) * torch.arange(half, dtype=torch.float32) / max(1, half)
+        )
+        self.register_buffer("freqs", freqs, persistent=False)
+        self.net = nn.Sequential(
+            nn.Linear(fourier_dim, out_dim),
+            nn.SiLU(),
+            nn.Linear(out_dim, out_dim),
+        )
+
+    def forward(self, x):
+        args = x.float()[:, None] * self.freqs[None]
+        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if self.fourier_dim % 2:
+            emb = F.pad(emb, (0, 1))
+        return self.net(emb)
+
+
+class ResBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, cond_dim):
+        super().__init__()
+        self.norm1 = JVPGroupNorm(gn_groups(in_ch), in_ch)
+        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
+        self.cond = nn.Linear(cond_dim, 2 * out_ch)
+        self.norm2 = JVPGroupNorm(gn_groups(out_ch), out_ch)
+        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
+        self.skip = nn.Identity() if in_ch == out_ch else nn.Conv2d(in_ch, out_ch, 1)
+        nn.init.zeros_(self.conv2.weight)
+        nn.init.zeros_(self.conv2.bias)
+
+    def forward(self, x, c):
+        h = self.conv1(F.silu(self.norm1(x)))
+        scale, shift = self.cond(F.silu(c)).chunk(2, dim=1)
+        h = self.norm2(h) * (1 + scale[:, :, None, None]) + shift[:, :, None, None]
+        return self.skip(x) + self.conv2(F.silu(h))
+
+
+class AttentionBlock(nn.Module):
+    """Static math attention path: forward-mode-AD safe and compile-friendly."""
+    def __init__(self, ch, max_heads=4):
+        super().__init__()
+        heads = min(max_heads, ch)
+        while ch % heads:
+            heads -= 1
+        self.heads = heads
+        self.norm = JVPGroupNorm(gn_groups(ch), ch)
+        self.qkv = nn.Conv1d(ch, 3 * ch, 1)
+        self.proj = nn.Conv1d(ch, ch, 1)
+        nn.init.zeros_(self.proj.weight)
+        nn.init.zeros_(self.proj.bias)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        L = h * w
+        z = self.norm(x).reshape(b, c, L)
+        q, k, v = self.qkv(z).chunk(3, dim=1)
+        d = c // self.heads
+        q = q.reshape(b, self.heads, d, L).transpose(-1, -2)
+        k = k.reshape(b, self.heads, d, L).transpose(-1, -2)
+        v = v.reshape(b, self.heads, d, L).transpose(-1, -2)
+        out = torch.softmax((q @ k.transpose(-1, -2)) * (d ** -0.5), dim=-1) @ v
+        out = self.proj(out.transpose(-1, -2).reshape(b, c, L)).reshape(b, c, h, w)
+        return x + out
+
+
+class Downsample(nn.Module):
+    def __init__(self, ch):
+        super().__init__()
+        self.op = nn.Conv2d(ch, ch, 3, stride=2, padding=1)
+
+    def forward(self, x):
+        return self.op(x)
+
+
+class Upsample(nn.Module):
+    def __init__(self, ch):
+        super().__init__()
+        self.op = nn.Conv2d(ch, ch, 3, padding=1)
+
+    def forward(self, x):
+        return self.op(F.interpolate(x, scale_factor=2.0, mode="nearest"))
+
+
+class ConditionalMultiTraceUNet(nn.Module):
+    """One shared conditional field u_theta(z,r,t,y)."""
+    def __init__(self, input_size=32, in_channels=3, num_classes=10,
+                 base_channels=84, channel_mult=(1, 2, 2, 4), num_res_blocks=1,
+                 cond_dim=356, attn_resolutions=(8,), fourier_dim=64):
+        super().__init__()
+        cfg = dict(image_size=input_size, num_classes=num_classes,
+                   base_channels=base_channels, channel_mult=channel_mult,
+                   num_res_blocks=num_res_blocks, cond_dim=cond_dim,
+                   attn_resolutions=attn_resolutions, fourier_dim=fourier_dim)
+        base = cfg["base_channels"]
+        mults = cfg["channel_mult"]
+        nrb = cfg["num_res_blocks"]
+        cd = cfg["cond_dim"]
+
+        self.attn_res = set(cfg["attn_resolutions"])
+        self.t_emb = ScalarEmbed(cd, cfg["fourier_dim"])
+        self.r_emb = ScalarEmbed(cd, cfg["fourier_dim"])
+        self.y_emb = nn.Embedding(cfg["num_classes"], cd)
+        nn.init.normal_(self.y_emb.weight, std=0.02)
+        self.cond_out = nn.Sequential(nn.SiLU(), nn.Linear(cd, cd))
+
+        self.in_conv = nn.Conv2d(in_channels, base, 3, padding=1)
+        self.down_levels = nn.ModuleList()
+        ch = base
+        res = cfg["image_size"]
+        skip_channels = [ch]
+
+        for li, m in enumerate(mults):
+            out_ch = base * m
+            blocks = nn.ModuleList()
+            atts = nn.ModuleList()
+            for _ in range(nrb):
+                blocks.append(ResBlock(ch, out_ch, cd))
+                ch = out_ch
+                atts.append(AttentionBlock(ch) if res in self.attn_res else nn.Identity())
+                skip_channels.append(ch)
+            down = Downsample(ch) if li != len(mults) - 1 else nn.Identity()
+            self.down_levels.append(nn.ModuleDict({"blocks": blocks, "atts": atts, "down": down}))
+            if li != len(mults) - 1:
+                skip_channels.append(ch)
+                res //= 2
+
+        self.mid1 = ResBlock(ch, ch, cd)
+        self.mid_attn = AttentionBlock(ch)
+        self.mid2 = ResBlock(ch, ch, cd)
+
+        self.up_levels = nn.ModuleList()
+        stack = list(skip_channels)
+        res = cfg["image_size"] // (2 ** (len(mults) - 1))
+        for rev_idx, m in enumerate(reversed(mults)):
+            li = len(mults) - 1 - rev_idx
+            out_ch = base * m
+            blocks = nn.ModuleList()
+            atts = nn.ModuleList()
+            for _ in range(nrb + 1):
+                skip_ch = stack.pop()
+                blocks.append(ResBlock(ch + skip_ch, out_ch, cd))
+                ch = out_ch
+                atts.append(AttentionBlock(ch) if res in self.attn_res else nn.Identity())
+            up = Upsample(ch) if li != 0 else nn.Identity()
+            self.up_levels.append(nn.ModuleDict({"blocks": blocks, "atts": atts, "up": up}))
+            if li != 0:
+                res *= 2
+
+        if stack:
+            raise RuntimeError(f"Unused build-time skips: {stack}")
+
+        self.out_norm = JVPGroupNorm(gn_groups(ch), ch)
+        self.u_head = nn.Conv2d(ch, in_channels, 3, padding=1)
+        nn.init.zeros_(self.u_head.weight)
+        nn.init.zeros_(self.u_head.bias)
+
+    def forward(self, x, t, r, y=None, use_flash_attention=False):
+        # Reference notebook uses (x,r,t,y); repository API uses (x,t,r,y).
+        if y is None:
+            raise ValueError("This U-Net requires class labels")
+        c = self.cond_out(self.t_emb(t) + self.r_emb(r) + self.y_emb(y))
+        hs = []
+        x = self.in_conv(x)
+        hs.append(x)
+
+        for level in self.down_levels:
+            for block, attn in zip(level["blocks"], level["atts"]):
+                x = attn(block(x, c))
+                hs.append(x)
+            if not isinstance(level["down"], nn.Identity):
+                x = level["down"](x)
+                hs.append(x)
+
+        x = self.mid2(self.mid_attn(self.mid1(x, c)), c)
+
+        for level in self.up_levels:
+            for block, attn in zip(level["blocks"], level["atts"]):
+                x = attn(block(torch.cat([x, hs.pop()], dim=1), c))
+            x = level["up"](x)
+
+        if hs:
+            raise RuntimeError(f"Unconsumed runtime skips: {len(hs)}")
+        return self.u_head(F.silu(self.out_norm(x)))

multi_trace_imf.py

@@ -0,0 +1,57 @@
+"""Equal-grid Multi-Trace objective and 1..K sampling from Untitled0.ipynb."""
+import torch
+from trace_imf import TraceIMF
+
+
+class MultiTraceIMF(TraceIMF):
+    def __init__(self, max_nfe=5, jvp_precision="bf16", channels_last=True, **kwargs):
+        super().__init__(**kwargs)
+        if max_nfe < 1 or jvp_precision not in ("fp32", "bf16"):
+            raise ValueError("Invalid max_nfe or JVP precision")
+        self.max_nfe = max_nfe
+        self.jvp_precision = jvp_precision
+        self.channels_last = channels_last
+
+    def draw_interval(self, device, generator=None, nfe=None):
+        # One K and j for the entire effective batch, including all microbatches.
+        k = int(torch.randint(1, self.max_nfe + 1, (), device=device, generator=generator)) if nfe is None else nfe
+        if not 1 <= k <= self.max_nfe:
+            raise ValueError("NFE outside the trained range")
+        j = int(torch.randint(k, (), device=device, generator=generator))
+        return k, j
+
+    def loss(self, model, x, labels=None, *, interval=None, nfe=None, generator=None,
+             t=None, r=None, noise=None, derivative_model=None):
+        if t is None:
+            k, j = interval or self.draw_interval(x.device, generator, nfe)
+            r = torch.full((len(x),), j / k, device=x.device)
+            t = r + torch.rand(len(x), device=x.device, generator=generator) / k
+        elif r is None:
+            raise ValueError("Explicit Multi-Trace t requires explicit r")
+        return super().loss(model, x, labels, derivative_model=derivative_model,
+                            t=t, r=r, noise=noise, generator=generator)
+
+    @torch.no_grad()
+    def sample(self, model, n_samples=None, labels=None, *, device=None, noise=None, nfe=1):
+        if not 1 <= nfe <= self.max_nfe:
+            raise ValueError("NFE outside the trained range")
+        device = device or next(model.parameters()).device
+        if labels is None:
+            raise ValueError("Class-conditioned sampling requires labels")
+        labels = labels.to(device)
+        z = (torch.randn(len(labels), self.channels, self.image_size, self.image_size, device=device)
+             if noise is None else noise.to(device).clone())
+        if len(z) != len(labels):
+            raise ValueError("Noise and labels must have the same batch size")
+        if self.channels_last:
+            z = z.contiguous(memory_format=torch.channels_last)
+        was_training = model.training
+        model.eval()
+        try:
+            for index in range(nfe, 0, -1):
+                t = torch.full((len(z),), index / nfe, device=device)
+                r = torch.full_like(t, (index - 1) / nfe)
+                z = z - model(z, t, r, y=labels, use_flash_attention=False) / nfe
+            return (z.clamp(-1, 1) + 1) * 0.5
+        finally:
+            model.train(was_training)

objectives.py

@@ -0,0 +1,15 @@
+"""Objective selection shared by training and the standalone inference bundle."""
+from trace_imf import TraceIMF
+from multi_trace_imf import MultiTraceIMF
+
+
+def build_objective(cfg):
+    common = dict(channels=cfg["model"]["in_channels"],
+                  image_size=cfg["model"]["input_size"],
+                  num_classes=cfg["model"]["num_classes"])
+    name = cfg.get("objective", "trace_imf")
+    if name == "multi_trace_imf":
+        return MultiTraceIMF(**common, **cfg["multi_trace_imf"])
+    if name == "trace_imf":
+        return TraceIMF(**common, **cfg["trace_imf"])
+    raise ValueError(f"Unknown objective: {name}")

requirements.txt

@@ -0,0 +1,5 @@
+torch>=2.2
+numpy>=1.24
+safetensors>=0.4
+pillow>=9
+huggingface_hub>=0.21

sample.py

@@ -0,0 +1,56 @@
+"""Sample the exported custom TraceDiT inference bundle without loading pickle."""
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+from safetensors.torch import load_file
+import torch
+
+from models import build_backbone
+from objectives import build_objective
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-dir", required=True)
+    parser.add_argument("--output", default="samples.png")
+    parser.add_argument("--n-per-class", type=int, default=4)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--weights", default="model.safetensors")
+    parser.add_argument("--nfe", type=int, default=1)
+    args = parser.parse_args()
+    if args.n_per_class <= 0:
+        parser.error("n-per-class must be positive")
+    folder = Path(args.model_dir)
+    cfg = json.loads((folder / "config.json").read_text())
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch.manual_seed(args.seed)
+    model = build_backbone(cfg).to(device)
+    model.load_state_dict(load_file(str(folder / args.weights), device=device), strict=True)
+    objective = build_objective(cfg)
+    if args.nfe != 1 and not hasattr(objective, "max_nfe"):
+        parser.error("This model supports only one NFE")
+    if objective.num_classes is None:
+        labels, columns = None, 10
+    else:
+        columns = objective.num_classes
+        labels = torch.arange(columns, device=device).repeat(args.n_per_class)
+    images = objective.sample(model, n_samples=columns * args.n_per_class,
+                              labels=labels, device=device,
+                              **({"nfe": args.nfe} if hasattr(objective, "max_nfe") else {}))
+    pixels = (images.permute(0, 2, 3, 1).cpu().numpy() * 255).round().astype(np.uint8)
+    rows = args.n_per_class
+    if pixels.shape[-1] == 1:
+        pixels = pixels[..., 0]
+    grid = np.concatenate([np.concatenate(pixels[row*columns:(row+1)*columns], axis=1)
+                           for row in range(rows)], axis=0)
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    Image.fromarray(grid).save(output)
+    print(f"{args.nfe}-NFE samples saved: {output}")
+
+
+if __name__ == "__main__":
+    main()

samples.jsonl

@@ -0,0 +1,45 @@
+{"step": 42000, "pixel_mean": 0.47372499108314514, "nfe": 1, "pixel_std": 0.23482201993465424, "between_samples_std": 0.2342541217803955}
+{"step": 42000, "pixel_mean": 0.4737403690814972, "nfe": 2, "pixel_std": 0.23318670690059662, "between_samples_std": 0.23249799013137817}
+{"step": 42000, "pixel_mean": 0.4733788073062897, "nfe": 3, "pixel_std": 0.23294048011302948, "between_samples_std": 0.23233386874198914}
+{"step": 42000, "pixel_mean": 0.47353851795196533, "nfe": 4, "pixel_std": 0.23217451572418213, "between_samples_std": 0.23162195086479187}
+{"step": 42000, "pixel_mean": 0.4734746813774109, "nfe": 5, "pixel_std": 0.2306039184331894, "between_samples_std": 0.23005102574825287}
+{"step": 43000, "pixel_mean": 0.4738962948322296, "nfe": 1, "pixel_std": 0.23467251658439636, "between_samples_std": 0.23411357402801514}
+{"step": 43000, "pixel_mean": 0.47399187088012695, "nfe": 2, "pixel_std": 0.23315978050231934, "between_samples_std": 0.23246732354164124}
+{"step": 43000, "pixel_mean": 0.4736575782299042, "nfe": 3, "pixel_std": 0.23255649209022522, "between_samples_std": 0.23195421695709229}
+{"step": 43000, "pixel_mean": 0.4738808572292328, "nfe": 4, "pixel_std": 0.23214948177337646, "between_samples_std": 0.23160535097122192}
+{"step": 43000, "pixel_mean": 0.4738691449165344, "nfe": 5, "pixel_std": 0.23027542233467102, "between_samples_std": 0.22973471879959106}
+{"step": 44000, "pixel_mean": 0.47436124086380005, "nfe": 1, "pixel_std": 0.23476190865039825, "between_samples_std": 0.23417595028877258}
+{"step": 44000, "pixel_mean": 0.474582701921463, "nfe": 2, "pixel_std": 0.23312358558177948, "between_samples_std": 0.23241935670375824}
+{"step": 44000, "pixel_mean": 0.47428470849990845, "nfe": 3, "pixel_std": 0.23243620991706848, "between_samples_std": 0.23181688785552979}
+{"step": 44000, "pixel_mean": 0.4745030701160431, "nfe": 4, "pixel_std": 0.23207718133926392, "between_samples_std": 0.2315332293510437}
+{"step": 44000, "pixel_mean": 0.4745500683784485, "nfe": 5, "pixel_std": 0.23035971820354462, "between_samples_std": 0.2298177182674408}
+{"step": 45000, "pixel_mean": 0.47422075271606445, "nfe": 1, "pixel_std": 0.23496539890766144, "between_samples_std": 0.23434244096279144}
+{"step": 45000, "pixel_mean": 0.4741981029510498, "nfe": 2, "pixel_std": 0.23349685966968536, "between_samples_std": 0.23275989294052124}
+{"step": 45000, "pixel_mean": 0.47374066710472107, "nfe": 3, "pixel_std": 0.23272140324115753, "between_samples_std": 0.23208104074001312}
+{"step": 45000, "pixel_mean": 0.4738740026950836, "nfe": 4, "pixel_std": 0.23238259553909302, "between_samples_std": 0.2318258285522461}
+{"step": 45000, "pixel_mean": 0.4740034341812134, "nfe": 5, "pixel_std": 0.2307807356119156, "between_samples_std": 0.23024126887321472}
+{"step": 46000, "pixel_mean": 0.4745643734931946, "nfe": 1, "pixel_std": 0.23507004976272583, "between_samples_std": 0.2344510555267334}
+{"step": 46000, "pixel_mean": 0.4744739830493927, "nfe": 2, "pixel_std": 0.23381182551383972, "between_samples_std": 0.2330787479877472}
+{"step": 46000, "pixel_mean": 0.47411832213401794, "nfe": 3, "pixel_std": 0.23263326287269592, "between_samples_std": 0.23199662566184998}
+{"step": 46000, "pixel_mean": 0.474417120218277, "nfe": 4, "pixel_std": 0.2324787825345993, "between_samples_std": 0.23192620277404785}
+{"step": 46000, "pixel_mean": 0.47444015741348267, "nfe": 5, "pixel_std": 0.2309519499540329, "between_samples_std": 0.23041489720344543}
+{"step": 47000, "pixel_mean": 0.47470709681510925, "nfe": 1, "pixel_std": 0.23535038530826569, "between_samples_std": 0.23473000526428223}
+{"step": 47000, "pixel_mean": 0.4745272994041443, "nfe": 2, "pixel_std": 0.2342585325241089, "between_samples_std": 0.23352518677711487}
+{"step": 47000, "pixel_mean": 0.47419074177742004, "nfe": 3, "pixel_std": 0.23334956169128418, "between_samples_std": 0.2327093929052353}
+{"step": 47000, "pixel_mean": 0.4744076132774353, "nfe": 4, "pixel_std": 0.23303648829460144, "between_samples_std": 0.23248785734176636}
+{"step": 47000, "pixel_mean": 0.47451314330101013, "nfe": 5, "pixel_std": 0.23162786662578583, "between_samples_std": 0.2310994565486908}
+{"step": 48000, "pixel_mean": 0.47464919090270996, "nfe": 1, "pixel_std": 0.2351643443107605, "between_samples_std": 0.2345452606678009}
+{"step": 48000, "pixel_mean": 0.47439342737197876, "nfe": 2, "pixel_std": 0.23403969407081604, "between_samples_std": 0.2333148717880249}
+{"step": 48000, "pixel_mean": 0.47428229451179504, "nfe": 3, "pixel_std": 0.23292940855026245, "between_samples_std": 0.2323130965232849}
+{"step": 48000, "pixel_mean": 0.47453099489212036, "nfe": 4, "pixel_std": 0.2325868308544159, "between_samples_std": 0.23205390572547913}
+{"step": 48000, "pixel_mean": 0.47459015250205994, "nfe": 5, "pixel_std": 0.23120051622390747, "between_samples_std": 0.23069819808006287}
+{"step": 49000, "pixel_mean": 0.4743174612522125, "nfe": 1, "pixel_std": 0.23494189977645874, "between_samples_std": 0.23430010676383972}
+{"step": 49000, "pixel_mean": 0.4741107225418091, "nfe": 2, "pixel_std": 0.23382453620433807, "between_samples_std": 0.23308944702148438}
+{"step": 49000, "pixel_mean": 0.47396624088287354, "nfe": 3, "pixel_std": 0.23279331624507904, "between_samples_std": 0.2321631908416748}
+{"step": 49000, "pixel_mean": 0.47409728169441223, "nfe": 4, "pixel_std": 0.2323838174343109, "between_samples_std": 0.2318393886089325}
+{"step": 49000, "pixel_mean": 0.47436952590942383, "nfe": 5, "pixel_std": 0.23112046718597412, "between_samples_std": 0.23060618340969086}
+{"step": 50000, "pixel_mean": 0.4741409122943878, "nfe": 1, "pixel_std": 0.23498037457466125, "between_samples_std": 0.23431378602981567}
+{"step": 50000, "pixel_mean": 0.47384196519851685, "nfe": 2, "pixel_std": 0.23412783443927765, "between_samples_std": 0.2333933562040329}
+{"step": 50000, "pixel_mean": 0.47364792227745056, "nfe": 3, "pixel_std": 0.2328539341688156, "between_samples_std": 0.23222188651561737}
+{"step": 50000, "pixel_mean": 0.47384998202323914, "nfe": 4, "pixel_std": 0.23271359503269196, "between_samples_std": 0.2321617305278778}
+{"step": 50000, "pixel_mean": 0.47402939200401306, "nfe": 5, "pixel_std": 0.23148764669895172, "between_samples_std": 0.23096564412117004}

source_manifest.json

@@ -0,0 +1,24 @@
+{
+  "repo_url": "https://github.com/namdeptraivcd/IMF.git",
+  "git_commit": "59d5164a2b20753b846568264ebb10c3e49beace",
+  "files": {
+    "train.py": "34f72a2bbba5341aeda9cd3fe021e18ed6bf1bbc73b823187a97ed83e03cc3aa",
+    "monitoring.py": "5315ee70585789a684162fbd4fa3a8f2479d5186d282d22a3fc314432b3100a9",
+    "hub.py": "ca23b68ef728ea2ae514918a09bb678d80df303fc97949593f2c6607cdc0fa03",
+    "sample.py": "90f1089547475b5ab1d8f5f8d160a33c639826eda95c6ca79ee2d3f527b24be1",
+    "models/dit.py": "843080c2702d596dd3b16e35b78501b168ca008d013d4394ee9c64c086361819",
+    "models/__init__.py": "1d9d5ec9cc26f0e76e570735a42e420314fdccc0a4aa113ddac642fd4d147c36",
+    "trace_imf.py": "43e2e8105541257376f648f3ee82a92cce8402530be59372f8d7ed0516501263",
+    "data.py": "a346a102fac83a40823501b68cf8055ac385ab3837cb6ef62b9212bdc7b9dffc",
+    "configs/__init__.py": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+    "configs/cifar10_22m.py": "1599fb09bab2ac556f46f434f82d22d61aa22901140679b4bfef62aa65245588",
+    "requirements-notebook.txt": "858b7587f8c9b4778b2f13955ce8a00ff1f9a8f1cbfabb0ce94c2e97c55cc5fc",
+    "models/unet.py": "57baed6e4ef61ca23aa0ddf62ec0c30b85e6d27f2504416269bb6d1245ee22df",
+    "multi_trace_imf.py": "a807809500be68da9d33f13fd528d987e61471161ad869f97c7411bb225fe5b8",
+    "objectives.py": "691d152ee1f0970fa9c279c596d09070c5a5c6618ac465bb462fac0188796f7d",
+    "ema.py": "1bc020387ec0020724d0a5c066180100f08204102fd79e804c3624203ebc4579",
+    "evaluation.py": "6f2d19263ea1b8c13c9fdc4678062adae99ea0773e0f03c8001f81a598f03238",
+    "configs/multitrace_cifar10_22m.py": "e0a8fc408170b618d88d19ccf71993af5f469e59c36c31dd38d920b57ee727bb",
+    "configs/multitrace_cifar10_15m.py": "96863dc436dce0accb6a2722b4398aa1a8a6ebf4e96386511767fe9b3c2db69d"
+  }
+}

tensorboard/segment_0041000_1791126419399065883/events.out.tfevents.1791126419.modal.498.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0419329574f8b6750d3cdfb4cd8da01f5d8d2c4ec603ae8cb6fed34a6bdb1527
+size 4065148

third_party/MeanFlow.LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Jiarui Hai
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

third_party/README.md

@@ -0,0 +1,18 @@
+# Nguồn code
+
+`models/dit.py`, `data.py` và cấu trúc training được kế thừa/adapt từ
+`../backbone_IMF`, nguồn [haidog-yaqub/MeanFlow](https://github.com/haidog-yaqub/MeanFlow),
+commit `97cc5182abe9d69e5dcb350268f603fbe31a5897`.
+MIT notice được giữ trong `MeanFlow.LICENSE`.
+
+Model giữ attention, adaLN-zero, embedding và unpatchify của nguồn; thay hai head
+bằng một head u, bỏ CFG embedding, cố định positional embedding, thay lớp timm
+bằng PyTorch tương đương và RMSNorm với epsilon 1e-6 để hỗ trợ PyTorch 2.2.
+Loss Trace-iMF được viết theo tài liệu idea của người dùng.
+
+`models/unet.py` giữ U-Net topology/shared head và `multi_trace_imf.py` giữ
+equal-grid interval schedule từ notebook `Untitled0.ipynb` do người dùng cung
+cấp. Bản 22M widen base44→84 và cond160→356; thứ tự API được chuẩn hóa thành
+`(z,t,r,y)`. GroupNorm nhận contiguous input để hỗ trợ forward AD trên PyTorch2.2.
+EMA/accumulation/resume/monitoring là phần tích hợp của IMF; không thực thi
+các cell Colab/unzip/Drive từ notebook tham chiếu.

trace_imf.py

@@ -0,0 +1,100 @@
+"""Two-trace quadratic objective from idea.pdf, sections 10 and 14."""
+import torch
+import torch.nn.functional as F
+
+
+class TraceIMF:
+    def __init__(self, channels=3, image_size=32, num_classes=10, lambda_diag=1.0):
+        if lambda_diag <= 0:
+            raise ValueError("lambda_diag must be positive to train both traces")
+        self.channels = channels
+        self.image_size = image_size
+        self.num_classes = num_classes
+        self.lambda_diag = lambda_diag
+
+    def loss(self, model, x, labels=None, *, derivative_model=None, t=None, noise=None,
+             r=None, generator=None):
+        """x is raw image data in [0,1]. Optional t/noise permit exact checks.
+
+        Use the unwrapped model for JVP under Accelerate/DDP. The single joint
+        grad-enabled call through model trains both edge and diagonal safely.
+        """
+        b = x.shape[0]
+        x = x.float() * 2 - 1
+        t = torch.rand(b, device=x.device, generator=generator) if t is None else t.to(x.device).float()
+        noise = torch.randn(x.shape, device=x.device, generator=generator) if noise is None else noise.to(x.device).float()
+        if t.shape != (b,) or noise.shape != x.shape:
+            raise ValueError("Expected t with shape (B,) and noise with the image shape")
+        if self.num_classes is not None and labels is None:
+            raise ValueError("Class-conditioned training requires labels")
+        if self.num_classes is None:
+            labels = None
+        t_image = t[:, None, None, None]
+        z = (1 - t_image) * x + t_image * noise
+        target = noise - x
+        r = torch.zeros_like(t) if r is None else r.to(x.device).float()
+        if r.shape != t.shape or (r < 0).any() or (t > 1).any() or (r > t).any():
+            raise ValueError("Require 0 <= r <= t <= 1, with r and t shaped (B,)")
+        if getattr(self, "channels_last", False):
+            z = z.contiguous(memory_format=torch.channels_last)
+
+        # Both branches use the SAME u head, evaluated at r=0 and r=t.
+        pair_labels = None if labels is None else torch.cat((labels, labels))
+        predictions = model(torch.cat((z, z)), torch.cat((t, t)),
+                            torch.cat((r, t)), y=pair_labels,
+                            use_flash_attention=True)
+        edge, diagonal = predictions.chunk(2)
+        jvp_model = model if derivative_model is None else derivative_model
+
+        # FP32 JVP: native attention supports forward AD; SDPA is used only in
+        # the grad-enabled forward. No parameter gradient through the JVP.
+        use_bf16 = getattr(self, "jvp_precision", "fp32") == "bf16" and x.device.type == "cuda"
+        with torch.no_grad(), torch.autocast(device_type=x.device.type, enabled=use_bf16,
+                                           dtype=torch.bfloat16):
+            def edge_fn(z_value, t_value, r_value):
+                # Accelerate may autocast the instance's forward even after
+                # unwrap_model. The class forward bypasses that AMP wrapper
+                # without mutating the grad-enabled training forward.
+                return type(jvp_model).forward(jvp_model, z_value, t_value, r_value,
+                                               y=labels, use_flash_attention=False)
+
+            _, derivative = torch.func.jvp(
+                edge_fn, (z, t, r),
+                (diagonal.detach().float(), torch.ones_like(t), torch.zeros_like(r)),
+            )
+
+        compound = edge.float() + (t - r)[:, None, None, None] * derivative.detach().float()
+        edge_loss = F.mse_loss(compound, target)
+        diag_loss = F.mse_loss(diagonal.float(), target)
+        loss = edge_loss + self.lambda_diag * diag_loss
+        return loss, {"edge_loss": edge_loss.detach(), "diag_loss": diag_loss.detach(),
+                      "jvp_rms": derivative.detach().square().mean().sqrt()}
+
+    @torch.no_grad()
+    def sample(self, model, n_samples=None, labels=None, *, device=None, noise=None):
+        """Exactly one evaluation at (t,r)=(1,0), then map back to [0,1]."""
+        was_training = model.training
+        model.eval()
+        try:
+            if device is None:
+                device = next(model.parameters()).device
+            if labels is not None:
+                labels = labels.to(device)
+                n_samples = len(labels)
+            if self.num_classes is not None and labels is None:
+                raise ValueError("Class-conditioned sampling requires labels")
+            if self.num_classes is None:
+                labels = None
+            if noise is None:
+                if n_samples is None or n_samples <= 0:
+                    raise ValueError("n_samples must be positive")
+                noise = torch.randn(n_samples, self.channels, self.image_size,
+                                    self.image_size, device=device)
+            else:
+                noise = noise.to(device)
+            t = torch.ones(noise.shape[0], device=device)
+            u = model(noise, t, torch.zeros_like(t), y=labels,
+                      use_flash_attention=True)
+            return ((noise - u).clamp(-1, 1) + 1) * 0.5
+        finally:
+            model.train(was_training)

training_config.json

@@ -0,0 +1,105 @@
+{
+  "name": "multitrace_unet_cifar10_15m",
+  "architecture": "ConditionalMultiTraceUNet",
+  "objective": "multi_trace_imf",
+  "seed": 42,
+  "dataset": "cifar10",
+  "data_root": "/mnt/imf-training/data/cifar10",
+  "image_size": 32,
+  "num_workers": 2,
+  "lr": 0.0001,
+  "min_lr_ratio": 0.1,
+  "betas": [
+    0.9,
+    0.99
+  ],
+  "weight_decay": 0.0,
+  "grad_clip": 1.0,
+  "mixed_precision": "bf16",
+  "channels_last": true,
+  "allow_tf32": true,
+  "fused_optimizer": true,
+  "preload_gpu": true,
+  "log_step": 50,
+  "sample_step": 1000,
+  "checkpoint_step": 1000,
+  "keep_last_checkpoints": 3,
+  "sample_nrow": 10,
+  "parameter_target": 15000000,
+  "parameter_tolerance": 0.01,
+  "model": {
+    "input_size": 32,
+    "in_channels": 3,
+    "num_classes": 10,
+    "base_channels": 69,
+    "channel_mult": [
+      1,
+      2,
+      2,
+      4
+    ],
+    "num_res_blocks": 1,
+    "cond_dim": 312,
+    "fourier_dim": 64,
+    "attn_resolutions": [
+      8
+    ]
+  },
+  "multi_trace_imf": {
+    "lambda_diag": 1.0,
+    "max_nfe": 5,
+    "jvp_precision": "bf16",
+    "channels_last": true
+  },
+  "monitoring": {
+    "enabled": true,
+    "tensorboard": true,
+    "save_best": true,
+    "validate_raw": true,
+    "plot_every": 500,
+    "validation_every": 1000,
+    "validation_batches": 5,
+    "sample_n_per_class": 2
+  },
+  "evaluation": {
+    "num_generated": 10000,
+    "batch_size": 128,
+    "nfes": [
+      1,
+      2,
+      3,
+      4,
+      5
+    ],
+    "seed": 2026
+  },
+  "batch_size": 1024,
+  "gradient_accumulation_steps": 1,
+  "n_steps": 50000,
+  "warmup_steps": 2500,
+  "ema_decay": 0.9998000100000001,
+  "source_manifest": {
+    "repo_url": "https://github.com/namdeptraivcd/IMF.git",
+    "git_commit": "59d5164a2b20753b846568264ebb10c3e49beace",
+    "files": {
+      "train.py": "34f72a2bbba5341aeda9cd3fe021e18ed6bf1bbc73b823187a97ed83e03cc3aa",
+      "monitoring.py": "5315ee70585789a684162fbd4fa3a8f2479d5186d282d22a3fc314432b3100a9",
+      "hub.py": "ca23b68ef728ea2ae514918a09bb678d80df303fc97949593f2c6607cdc0fa03",
+      "sample.py": "90f1089547475b5ab1d8f5f8d160a33c639826eda95c6ca79ee2d3f527b24be1",
+      "models/dit.py": "843080c2702d596dd3b16e35b78501b168ca008d013d4394ee9c64c086361819",
+      "models/__init__.py": "1d9d5ec9cc26f0e76e570735a42e420314fdccc0a4aa113ddac642fd4d147c36",
+      "trace_imf.py": "43e2e8105541257376f648f3ee82a92cce8402530be59372f8d7ed0516501263",
+      "data.py": "a346a102fac83a40823501b68cf8055ac385ab3837cb6ef62b9212bdc7b9dffc",
+      "configs/__init__.py": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+      "configs/cifar10_22m.py": "1599fb09bab2ac556f46f434f82d22d61aa22901140679b4bfef62aa65245588",
+      "requirements-notebook.txt": "858b7587f8c9b4778b2f13955ce8a00ff1f9a8f1cbfabb0ce94c2e97c55cc5fc",
+      "models/unet.py": "57baed6e4ef61ca23aa0ddf62ec0c30b85e6d27f2504416269bb6d1245ee22df",
+      "multi_trace_imf.py": "a807809500be68da9d33f13fd528d987e61471161ad869f97c7411bb225fe5b8",
+      "objectives.py": "691d152ee1f0970fa9c279c596d09070c5a5c6618ac465bb462fac0188796f7d",
+      "ema.py": "1bc020387ec0020724d0a5c066180100f08204102fd79e804c3624203ebc4579",
+      "evaluation.py": "6f2d19263ea1b8c13c9fdc4678062adae99ea0773e0f03c8001f81a598f03238",
+      "configs/multitrace_cifar10_22m.py": "e0a8fc408170b618d88d19ccf71993af5f469e59c36c31dd38d920b57ee727bb",
+      "configs/multitrace_cifar10_15m.py": "96863dc436dce0accb6a2722b4398aa1a8a6ebf4e96386511767fe9b3c2db69d"
+    }
+  }
+}

training_metrics.json

@@ -0,0 +1,198 @@
+{
+  "step": 50000,
+  "lr": 1.0000000098422651e-05,
+  "window_steps": 50,
+  "elapsed_seconds": 5178.840372322999,
+  "steps_per_second": 1.7378407815190096,
+  "eta_seconds": 0.0,
+  "images_seen": 51200000,
+  "images_per_second": 1779.5489602754658,
+  "effective_batch_size": 1024,
+  "loss": 0.3548988401889801,
+  "edge_loss": 0.18379321694374084,
+  "diag_loss": 0.17110557854175568,
+  "jvp_rms": 1.8925390243530273,
+  "grad_norm": 0.21770034730434418,
+  "grad_norm_postclip": 0.21770034730434418,
+  "clip_fraction": 0.0,
+  "ema_updates": 50000,
+  "ema_decay": 0.9998000100000001,
+  "loss_nfe_3": 0.35597619265317915,
+  "updates_nfe_3": 10,
+  "loss_nfe_4": 0.37776655620998806,
+  "updates_nfe_4": 9,
+  "loss_nfe_5": 0.3045877367258072,
+  "updates_nfe_5": 8,
+  "loss_nfe_2": 0.3413205325603485,
+  "updates_nfe_2": 15,
+  "loss_nfe_1": 0.403596218675375,
+  "updates_nfe_1": 8,
+  "gpu_allocated_gb": 0.8975429534912109,
+  "gpu_reserved_gb": 43.775390625,
+  "gpu_peak_allocated_gb": 37.954071044921875,
+  "gradient_groups": {
+    "t_emb": {
+      "weight_norm": 16.515023530024582,
+      "grad_norm": 0.11432138125518912,
+      "grad_rms": 0.00033289269578588913,
+      "grad_to_weight": 0.00692226572050297,
+      "update_to_weight": 5.734569434715802e-05,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    },
+    "r_emb": {
+      "weight_norm": 16.262888330958752,
+      "grad_norm": 0.039618767773493664,
+      "grad_rms": 0.00011536598196266779,
+      "grad_to_weight": 0.002436145841207901,
+      "update_to_weight": 4.733992422527738e-05,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    },
+    "y_emb": {
+      "weight_norm": 5.067645060453645,
+      "grad_norm": 0.0034476322407670404,
+      "grad_rms": 6.172251819577942e-05,
+      "grad_to_weight": 0.000680322358736469,
+      "update_to_weight": 2.5612938476024494e-05,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    },
+    "cond_out": {
+      "weight_norm": 12.87244278780043,
+      "grad_norm": 0.03697475371479452,
+      "grad_rms": 0.00011831936333622446,
+      "grad_to_weight": 0.002872396042019042,
+      "update_to_weight": 6.601609987986573e-05,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    },
+    "in_conv": {
+      "weight_norm": 5.039121547619157,
+      "grad_norm": 0.02010702286093248,
+      "grad_rms": 0.000457450609587797,
+      "grad_to_weight": 0.0039901841364458615,
+      "update_to_weight": 3.290046528251253e-05,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    },
+    "down_levels.0": {
+      "weight_norm": 16.64608276842342,
+      "grad_norm": 0.09815406752986451,
+      "grad_rms": 0.0002365165120239306,
+      "grad_to_weight": 0.0058965264618326095,
+      "update_to_weight": 8.954140929857495e-05,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    },
+    "down_levels.1": {
+      "weight_norm": 24.577499197373093,
+      "grad_norm": 0.01800084380436817,
+      "grad_rms": 2.4834869479276555e-05,
+      "grad_to_weight": 0.0007324115305552383,
+      "update_to_weight": 9.463347569559213e-05,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    },
+    "down_levels.2": {
+      "weight_norm": 30.83386603917639,
+      "grad_norm": 0.008038267314234277,
+      "grad_rms": 9.758269257084616e-06,
+      "grad_to_weight": 0.0002606960575109572,
+      "update_to_weight": 6.642526625302697e-05,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    },
+    "down_levels.3": {
+      "weight_norm": 34.63907268454675,
+      "grad_norm": 0.003863454481162671,
+      "grad_rms": 3.468231732491786e-06,
+      "grad_to_weight": 0.00011153458166581472,
+      "update_to_weight": 7.572530100937947e-05,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    },
+    "mid1": {
+      "weight_norm": 37.09438513040922,
+      "grad_norm": 0.0032076774317542586,
+      "grad_rms": 2.580133241292384e-06,
+      "grad_to_weight": 8.647339537984874e-05,
+      "update_to_weight": 7.81818486024589e-05,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    },
+    "mid_attn": {
+      "weight_norm": 24.673390202443425,
+      "grad_norm": 0.0011791270861310043,
+      "grad_rms": 2.1303187281934394e-06,
+      "grad_to_weight": 4.778942319869097e-05,
+      "update_to_weight": 5.262799180821532e-05,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    },
+    "mid2": {
+      "weight_norm": 37.37197710653749,
+      "grad_norm": 0.003229066174170447,
+      "grad_rms": 2.597337528965184e-06,
+      "grad_to_weight": 8.640340769141662e-05,
+      "update_to_weight": 7.984646921565486e-05,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    },
+    "up_levels.0": {
+      "weight_norm": 63.52169789540913,
+      "grad_norm": 0.008797741452653435,
+      "grad_rms": 3.905892380033334e-06,
+      "grad_to_weight": 0.0001384997842334008,
+      "update_to_weight": 9.197308783625432e-05,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    },
+    "up_levels.1": {
+      "weight_norm": 50.788316815976486,
+      "grad_norm": 0.022728171192384107,
+      "grad_rms": 1.695910240002863e-05,
+      "grad_to_weight": 0.00044750786435266354,
+      "update_to_weight": 8.085612339796532e-05,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    },
+    "up_levels.2": {
+      "weight_norm": 40.116334734120464,
+      "grad_norm": 0.038615167286647545,
+      "grad_rms": 3.316118164288409e-05,
+      "grad_to_weight": 0.0009625796459865482,
+      "update_to_weight": 9.42659917428506e-05,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    },
+    "up_levels.3": {
+      "weight_norm": 28.206099436294913,
+      "grad_norm": 0.10106223952213236,
+      "grad_rms": 0.00015754171001242578,
+      "grad_to_weight": 0.0035829923861109263,
+      "update_to_weight": 8.058384138680281e-05,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    },
+    "out_norm": {
+      "weight_norm": 8.418451806035767,
+      "grad_norm": 0.00825520231395145,
+      "grad_rms": 0.0007027294882485905,
+      "grad_to_weight": 0.0009806081336752117,
+      "update_to_weight": 5.2765190801682055e-06,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    },
+    "u_head": {
+      "weight_norm": 1.3251545617986,
+      "grad_norm": 0.10927559814558875,
+      "grad_rms": 0.002529690400215243,
+      "grad_to_weight": 0.08246253025554366,
+      "update_to_weight": 0.00010683296531524654,
+      "missing_grad_tensors": 0,
+      "zero_grad_tensors": 0
+    }
+  },
+  "gradient_snapshot_step": 50000
+}

validation.jsonl

@@ -0,0 +1,10 @@
+{"step": 41000, "weights_source": "ema", "loss": 0.3725855791568756, "edge_loss": 0.19200483143329622, "diag_loss": 0.18058074414730071, "jvp_rms": 2.078083175420761, "examples": 5120, "loss_nfe_1": 0.3927511692047119, "edge_loss_nfe_1": 0.2175933986902237, "diag_loss_nfe_1": 0.17515776455402374, "jvp_rms_nfe_1": 2.4403680324554444, "loss_nfe_2": 0.369793975353241, "edge_loss_nfe_2": 0.18860338628292084, "diag_loss_nfe_2": 0.1811905950307846, "jvp_rms_nfe_2": 2.2941908717155455, "loss_nfe_3": 0.3564019322395325, "edge_loss_nfe_3": 0.17990776002407075, "diag_loss_nfe_3": 0.17649417072534562, "jvp_rms_nfe_3": 1.885271716117859, "loss_nfe_4": 0.3897887945175171, "edge_loss_nfe_4": 0.19601401388645173, "diag_loss_nfe_4": 0.1937747687101364, "jvp_rms_nfe_4": 2.270489048957825, "loss_nfe_5": 0.3541920244693756, "edge_loss_nfe_5": 0.17790559828281402, "diag_loss_nfe_5": 0.17628642171621323, "jvp_rms_nfe_5": 1.500096207857132, "raw_loss": 0.379677836894989, "raw_edge_loss": 0.19833322674036027, "raw_diag_loss": 0.1813446119427681, "raw_jvp_rms": 2.2722322368621826, "raw_loss_nfe_1": 0.4198586106300354, "raw_edge_loss_nfe_1": 0.24385497570037842, "raw_diag_loss_nfe_1": 0.17600363194942475, "raw_jvp_rms_nfe_1": 2.645507574081421, "raw_loss_nfe_2": 0.37339203953742983, "raw_edge_loss_nfe_2": 0.19148991107940674, "raw_diag_loss_nfe_2": 0.18190212845802306, "raw_jvp_rms_nfe_2": 2.5214332699775697, "raw_loss_nfe_3": 0.35822181701660155, "raw_edge_loss_nfe_3": 0.18091285228729248, "raw_diag_loss_nfe_3": 0.1773089662194252, "raw_jvp_rms_nfe_3": 2.062240254878998, "raw_loss_nfe_4": 0.3911686480045319, "raw_edge_loss_nfe_4": 0.19675154983997345, "raw_diag_loss_nfe_4": 0.1944171041250229, "raw_jvp_rms_nfe_4": 2.484368699789047, "raw_loss_nfe_5": 0.35574806928634645, "raw_edge_loss_nfe_5": 0.17865684479475022, "raw_diag_loss_nfe_5": 0.17709122896194457, "raw_jvp_rms_nfe_5": 1.6476113855838777}
+{"step": 42000, "weights_source": "ema", "loss": 0.3723788583278656, "edge_loss": 0.19195123821496965, "diag_loss": 0.18042762011289598, "jvp_rms": 2.0993320310115813, "examples": 5120, "loss_nfe_1": 0.3929444789886475, "edge_loss_nfe_1": 0.21792266964912416, "diag_loss_nfe_1": 0.17502180635929107, "jvp_rms_nfe_1": 2.471006441116333, "loss_nfe_2": 0.36948872208595274, "edge_loss_nfe_2": 0.18845009803771973, "diag_loss_nfe_2": 0.18103862702846527, "jvp_rms_nfe_2": 2.3163293421268465, "loss_nfe_3": 0.35611130595207213, "edge_loss_nfe_3": 0.1797650545835495, "diag_loss_nfe_3": 0.17634625285863875, "jvp_rms_nfe_3": 1.9045085847377776, "loss_nfe_4": 0.3894401788711548, "edge_loss_nfe_4": 0.19585056602954865, "diag_loss_nfe_4": 0.19358961880207062, "jvp_rms_nfe_4": 2.2923266768455504, "loss_nfe_5": 0.3539096057415009, "edge_loss_nfe_5": 0.17776780277490617, "diag_loss_nfe_5": 0.1761417955160141, "jvp_rms_nfe_5": 1.5124891102313995, "raw_loss": 0.37814890921115873, "raw_edge_loss": 0.19729300439357758, "raw_diag_loss": 0.1808559051156044, "raw_jvp_rms": 2.621889500617981, "raw_loss_nfe_1": 0.4127004086971283, "raw_edge_loss_nfe_1": 0.23728247880935668, "raw_diag_loss_nfe_1": 0.17541792690753938, "raw_jvp_rms_nfe_1": 3.186359691619873, "raw_loss_nfe_2": 0.37278689742088317, "raw_edge_loss_nfe_2": 0.19128688871860505, "raw_diag_loss_nfe_2": 0.18150001168251037, "raw_jvp_rms_nfe_2": 2.9029066443443297, "raw_loss_nfe_3": 0.3582378625869751, "raw_edge_loss_nfe_3": 0.18145101070404052, "raw_diag_loss_nfe_3": 0.1767868474125862, "raw_jvp_rms_nfe_3": 2.349283879995346, "raw_loss_nfe_4": 0.3917541980743408, "raw_edge_loss_nfe_4": 0.19763608276844025, "raw_diag_loss_nfe_4": 0.1941181242465973, "raw_jvp_rms_nfe_4": 2.840165650844574, "raw_loss_nfe_5": 0.3552651792764664, "raw_edge_loss_nfe_5": 0.17880856096744538, "raw_diag_loss_nfe_5": 0.17645661532878876, "raw_jvp_rms_nfe_5": 1.830731636285782}
+{"step": 43000, "weights_source": "ema", "loss": 0.37205516755580903, "edge_loss": 0.1917748808860779, "diag_loss": 0.18028028905391694, "jvp_rms": 2.125342252254486, "examples": 5120, "loss_nfe_1": 0.3925249338150024, "edge_loss_nfe_1": 0.21763605773448944, "diag_loss_nfe_1": 0.17488887906074524, "jvp_rms_nfe_1": 2.5001804351806642, "loss_nfe_2": 0.3691649556159973, "edge_loss_nfe_2": 0.1882774442434311, "diag_loss_nfe_2": 0.18088751137256623, "jvp_rms_nfe_2": 2.3510846734046935, "loss_nfe_3": 0.35585449934005736, "edge_loss_nfe_3": 0.17965261936187743, "diag_loss_nfe_3": 0.1762018769979477, "jvp_rms_nfe_3": 1.92358438372612, "loss_nfe_4": 0.3890950262546539, "edge_loss_nfe_4": 0.19568087458610534, "diag_loss_nfe_4": 0.1934141606092453, "jvp_rms_nfe_4": 2.3233995378017425, "loss_nfe_5": 0.35363642275333407, "edge_loss_nfe_5": 0.17762740850448608, "diag_loss_nfe_5": 0.1760090172290802, "jvp_rms_nfe_5": 1.5284622311592102, "raw_loss": 0.37374089658260345, "raw_edge_loss": 0.1935477900505066, "raw_diag_loss": 0.1801931056380272, "raw_jvp_rms": 2.299976586103439, "raw_loss_nfe_1": 0.39992823004722594, "raw_edge_loss_nfe_1": 0.2250307857990265, "raw_diag_loss_nfe_1": 0.17489744424819947, "raw_jvp_rms_nfe_1": 2.7114959239959715, "raw_loss_nfe_2": 0.36982461214065554, "raw_edge_loss_nfe_2": 0.18907322287559508, "raw_diag_loss_nfe_2": 0.18075139224529266, "raw_jvp_rms_nfe_2": 2.5491456270217894, "raw_loss_nfe_3": 0.3560143828392029, "raw_edge_loss_nfe_3": 0.17990127801895142, "raw_diag_loss_nfe_3": 0.17611309438943862, "raw_jvp_rms_nfe_3": 2.0785558223724365, "raw_loss_nfe_4": 0.3891355276107788, "raw_edge_loss_nfe_4": 0.19592745006084442, "raw_diag_loss_nfe_4": 0.19320808351039886, "raw_jvp_rms_nfe_4": 2.514833205938339, "raw_loss_nfe_5": 0.3538017302751541, "raw_edge_loss_nfe_5": 0.17780621349811554, "raw_diag_loss_nfe_5": 0.17599551379680634, "raw_jvp_rms_nfe_5": 1.6458523511886596}
+{"step": 44000, "weights_source": "ema", "loss": 0.3718108052015304, "edge_loss": 0.19166429877281188, "diag_loss": 0.1801465055346489, "jvp_rms": 2.1521597826480865, "examples": 5120, "loss_nfe_1": 0.3922906398773193, "edge_loss_nfe_1": 0.2175253540277481, "diag_loss_nfe_1": 0.17476528584957124, "jvp_rms_nfe_1": 2.536211538314819, "loss_nfe_2": 0.3689300060272217, "edge_loss_nfe_2": 0.18817995488643646, "diag_loss_nfe_2": 0.18075005412101747, "jvp_rms_nfe_2": 2.3827747225761415, "loss_nfe_3": 0.35562258362770083, "edge_loss_nfe_3": 0.17955487966537476, "diag_loss_nfe_3": 0.17606770247220993, "jvp_rms_nfe_3": 1.947074192762375, "loss_nfe_4": 0.38880528807640075, "edge_loss_nfe_4": 0.19554311633110047, "diag_loss_nfe_4": 0.19326217472553253, "jvp_rms_nfe_4": 2.3517152607440948, "loss_nfe_5": 0.3534055083990097, "edge_loss_nfe_5": 0.17751818895339966, "diag_loss_nfe_5": 0.17588731050491332, "jvp_rms_nfe_5": 1.5430231988430023, "raw_loss": 0.37199350357055666, "raw_edge_loss": 0.19197442293167116, "raw_diag_loss": 0.18001908123493196, "raw_jvp_rms": 2.215284506082535, "raw_loss_nfe_1": 0.3944181323051453, "raw_edge_loss_nfe_1": 0.2197323888540268, "raw_diag_loss_nfe_1": 0.1746857464313507, "raw_jvp_rms_nfe_1": 2.620998001098633, "raw_loss_nfe_2": 0.3687027633190155, "raw_edge_loss_nfe_2": 0.18810227811336516, "raw_diag_loss_nfe_2": 0.18060048818588256, "raw_jvp_rms_nfe_2": 2.459484040737152, "raw_loss_nfe_3": 0.35545286536216736, "raw_edge_loss_nfe_3": 0.17949930727481841, "raw_diag_loss_nfe_3": 0.17595355361700057, "raw_jvp_rms_nfe_3": 2.002295124530792, "raw_loss_nfe_4": 0.38823755383491515, "raw_edge_loss_nfe_4": 0.19519887268543243, "raw_diag_loss_nfe_4": 0.1930386781692505, "raw_jvp_rms_nfe_4": 2.4217930257320406, "raw_loss_nfe_5": 0.35315620303153994, "raw_edge_loss_nfe_5": 0.1773392677307129, "raw_diag_loss_nfe_5": 0.17581693977117538, "raw_jvp_rms_nfe_5": 1.5718523383140564}
+{"step": 45000, "weights_source": "ema", "loss": 0.37168434262275696, "edge_loss": 0.19165305346250533, "diag_loss": 0.18003128856420517, "jvp_rms": 2.1736205995082853, "examples": 5120, "loss_nfe_1": 0.3925355851650238, "edge_loss_nfe_1": 0.21787453591823577, "diag_loss_nfe_1": 0.17466104328632354, "jvp_rms_nfe_1": 2.5622494220733643, "loss_nfe_2": 0.3687441825866699, "edge_loss_nfe_2": 0.18811105489730834, "diag_loss_nfe_2": 0.18063313663005828, "jvp_rms_nfe_2": 2.406940096616745, "loss_nfe_3": 0.35541791915893556, "edge_loss_nfe_3": 0.1794615715742111, "diag_loss_nfe_3": 0.17595634311437608, "jvp_rms_nfe_3": 1.9702612817287446, "loss_nfe_4": 0.3885439932346344, "edge_loss_nfe_4": 0.19541764855384827, "diag_loss_nfe_4": 0.1931263417005539, "jvp_rms_nfe_4": 2.373984634876251, "loss_nfe_5": 0.35318003296852113, "edge_loss_nfe_5": 0.17740045636892318, "diag_loss_nfe_5": 0.17577957808971406, "jvp_rms_nfe_5": 1.5546675622463226, "raw_loss": 0.3740403318405151, "raw_edge_loss": 0.1942245441675186, "raw_diag_loss": 0.17981578558683395, "raw_jvp_rms": 2.3283275413513183, "raw_loss_nfe_1": 0.4041799962520599, "raw_edge_loss_nfe_1": 0.2296959787607193, "raw_diag_loss_nfe_1": 0.1744840145111084, "raw_jvp_rms_nfe_1": 2.760013055801392, "raw_loss_nfe_2": 0.36933073997497556, "raw_edge_loss_nfe_2": 0.18892420828342438, "raw_diag_loss_nfe_2": 0.18040653467178344, "raw_jvp_rms_nfe_2": 2.587478423118591, "raw_loss_nfe_3": 0.3554180383682251, "raw_edge_loss_nfe_3": 0.17966652810573577, "raw_diag_loss_nfe_3": 0.1757515013217926, "raw_jvp_rms_nfe_3": 2.1034037828445435, "raw_loss_nfe_4": 0.38829761147499087, "raw_edge_loss_nfe_4": 0.1954770267009735, "raw_diag_loss_nfe_4": 0.19282058477401734, "raw_jvp_rms_nfe_4": 2.5383615255355836, "raw_loss_nfe_5": 0.35297527313232424, "raw_edge_loss_nfe_5": 0.1773589789867401, "raw_diag_loss_nfe_5": 0.175616292655468, "raw_jvp_rms_nfe_5": 1.652380919456482}
+{"step": 46000, "weights_source": "ema", "loss": 0.3715553200244904, "edge_loss": 0.19162928521633146, "diag_loss": 0.1799260401725769, "jvp_rms": 2.201610367298126, "examples": 5120, "loss_nfe_1": 0.3926394283771515, "edge_loss_nfe_1": 0.2180699735879898, "diag_loss_nfe_1": 0.17456946074962615, "jvp_rms_nfe_1": 2.6001721382141114, "loss_nfe_2": 0.3686067283153534, "edge_loss_nfe_2": 0.18808059692382811, "diag_loss_nfe_2": 0.18052614033222197, "jvp_rms_nfe_2": 2.439925366640091, "loss_nfe_3": 0.35523309707641604, "edge_loss_nfe_3": 0.1793862372636795, "diag_loss_nfe_3": 0.17584685534238814, "jvp_rms_nfe_3": 1.992531567811966, "loss_nfe_4": 0.38829939961433413, "edge_loss_nfe_4": 0.19529785215854645, "diag_loss_nfe_4": 0.19300155341625214, "jvp_rms_nfe_4": 2.4034175872802734, "loss_nfe_5": 0.35299794673919677, "edge_loss_nfe_5": 0.17731176614761351, "diag_loss_nfe_5": 0.17568619102239608, "jvp_rms_nfe_5": 1.5720051765441894, "raw_loss": 0.37294554710388184, "raw_edge_loss": 0.19305831611156463, "raw_diag_loss": 0.17988722890615463, "raw_jvp_rms": 2.5546124279499054, "raw_loss_nfe_1": 0.39817399382591245, "raw_edge_loss_nfe_1": 0.2236364871263504, "raw_diag_loss_nfe_1": 0.17453750371932983, "raw_jvp_rms_nfe_1": 3.0601311683654786, "raw_loss_nfe_2": 0.36934937834739684, "raw_edge_loss_nfe_2": 0.18887795507907867, "raw_diag_loss_nfe_2": 0.1804714173078537, "raw_jvp_rms_nfe_2": 2.8419912278652193, "raw_loss_nfe_3": 0.35562482476234436, "raw_edge_loss_nfe_3": 0.17977170050144195, "raw_diag_loss_nfe_3": 0.17585312575101852, "raw_jvp_rms_nfe_3": 2.303293800354004, "raw_loss_nfe_4": 0.3885082483291626, "raw_edge_loss_nfe_4": 0.19556552171707153, "raw_diag_loss_nfe_4": 0.1929427295923233, "raw_jvp_rms_nfe_4": 2.7833980321884155, "raw_loss_nfe_5": 0.3530712902545929, "raw_edge_loss_nfe_5": 0.1774399161338806, "raw_diag_loss_nfe_5": 0.1756313681602478, "raw_jvp_rms_nfe_5": 1.78424791097641}
+{"step": 47000, "weights_source": "ema", "loss": 0.3713958901166916, "edge_loss": 0.1915695336461067, "diag_loss": 0.1798263558745384, "jvp_rms": 2.2319657707214353, "examples": 5120, "loss_nfe_1": 0.3926237761974335, "edge_loss_nfe_1": 0.2181445300579071, "diag_loss_nfe_1": 0.1744792491197586, "jvp_rms_nfe_1": 2.642152261734009, "loss_nfe_2": 0.36840869188308717, "edge_loss_nfe_2": 0.18798398077487946, "diag_loss_nfe_2": 0.18042470812797545, "jvp_rms_nfe_2": 2.475356990098953, "loss_nfe_3": 0.3550576686859131, "edge_loss_nfe_3": 0.17930547297000884, "diag_loss_nfe_3": 0.175752192735672, "jvp_rms_nfe_3": 2.0165530383586883, "loss_nfe_4": 0.3880728900432587, "edge_loss_nfe_4": 0.19519352316856384, "diag_loss_nfe_4": 0.19287936985492707, "jvp_rms_nfe_4": 2.435269385576248, "loss_nfe_5": 0.35281642377376554, "edge_loss_nfe_5": 0.17722016125917434, "diag_loss_nfe_5": 0.17559625953435898, "jvp_rms_nfe_5": 1.5904971778392791, "raw_loss": 0.3728845959901809, "raw_edge_loss": 0.19330160319805145, "raw_diag_loss": 0.17958299428224564, "raw_jvp_rms": 2.488507117033005, "raw_loss_nfe_1": 0.4004963099956512, "raw_edge_loss_nfe_1": 0.226207235455513, "raw_diag_loss_nfe_1": 0.1742890805006027, "raw_jvp_rms_nfe_1": 2.9745230674743652, "raw_loss_nfe_2": 0.36872340440750123, "raw_edge_loss_nfe_2": 0.18855189383029938, "raw_diag_loss_nfe_2": 0.18017151057720185, "raw_jvp_rms_nfe_2": 2.7699862122535706, "raw_loss_nfe_3": 0.3549445480108261, "raw_edge_loss_nfe_3": 0.17942589223384858, "raw_diag_loss_nfe_3": 0.17551865577697753, "raw_jvp_rms_nfe_3": 2.246046853065491, "raw_loss_nfe_4": 0.3877571880817413, "raw_edge_loss_nfe_4": 0.19519511461257935, "raw_diag_loss_nfe_4": 0.192562073469162, "raw_jvp_rms_nfe_4": 2.705729341506958, "raw_loss_nfe_5": 0.3525015294551849, "raw_edge_loss_nfe_5": 0.17712787985801698, "raw_diag_loss_nfe_5": 0.1753736510872841, "raw_jvp_rms_nfe_5": 1.7462501108646393}
+{"step": 48000, "weights_source": "ema", "loss": 0.3712373948097229, "edge_loss": 0.19149310767650604, "diag_loss": 0.17974428683519364, "jvp_rms": 2.250151746273041, "examples": 5120, "loss_nfe_1": 0.39254979491233827, "edge_loss_nfe_1": 0.21814591586589813, "diag_loss_nfe_1": 0.17440388202667237, "jvp_rms_nfe_1": 2.665694570541382, "loss_nfe_2": 0.3681997716426849, "edge_loss_nfe_2": 0.18785875737667085, "diag_loss_nfe_2": 0.1803410142660141, "jvp_rms_nfe_2": 2.495949363708496, "loss_nfe_3": 0.3549005001783371, "edge_loss_nfe_3": 0.17922913730144502, "diag_loss_nfe_3": 0.17567135840654374, "jvp_rms_nfe_3": 2.0347792744636535, "loss_nfe_4": 0.3878713846206665, "edge_loss_nfe_4": 0.19508881270885467, "diag_loss_nfe_4": 0.19278257191181183, "jvp_rms_nfe_4": 2.4547479271888735, "loss_nfe_5": 0.35266552269458773, "edge_loss_nfe_5": 0.17714291512966157, "diag_loss_nfe_5": 0.17552260756492616, "jvp_rms_nfe_5": 1.5995875954627992, "raw_loss": 0.37239629626274107, "raw_edge_loss": 0.19247002869844437, "raw_diag_loss": 0.17992626249790195, "raw_jvp_rms": 2.1689469707012177, "raw_loss_nfe_1": 0.396082466840744, "raw_edge_loss_nfe_1": 0.2214716374874115, "raw_diag_loss_nfe_1": 0.17461082339286804, "raw_jvp_rms_nfe_1": 2.5685171604156496, "raw_loss_nfe_2": 0.36995372772216795, "raw_edge_loss_nfe_2": 0.18943111300468446, "raw_diag_loss_nfe_2": 0.1805226117372513, "raw_jvp_rms_nfe_2": 2.4097091972827913, "raw_loss_nfe_3": 0.35518380999565125, "raw_edge_loss_nfe_3": 0.17940330505371094, "raw_diag_loss_nfe_3": 0.17578050196170808, "raw_jvp_rms_nfe_3": 1.9646742224693299, "raw_loss_nfe_4": 0.3877831220626831, "raw_edge_loss_nfe_4": 0.19488065540790558, "raw_diag_loss_nfe_4": 0.19290245771408082, "raw_jvp_rms_nfe_4": 2.366553395986557, "raw_loss_nfe_5": 0.3529783546924591, "raw_edge_loss_nfe_5": 0.17716343253850936, "raw_diag_loss_nfe_5": 0.17581491768360138, "raw_jvp_rms_nfe_5": 1.5352808773517608}
+{"step": 49000, "weights_source": "ema", "loss": 0.3710547721385956, "edge_loss": 0.19137788712978362, "diag_loss": 0.17967688530683518, "jvp_rms": 2.26550262093544, "examples": 5120, "loss_nfe_1": 0.39221864342689516, "edge_loss_nfe_1": 0.21786801517009735, "diag_loss_nfe_1": 0.17435062825679778, "jvp_rms_nfe_1": 2.686102199554443, "loss_nfe_2": 0.3680577754974365, "edge_loss_nfe_2": 0.1877869725227356, "diag_loss_nfe_2": 0.18027080297470094, "jvp_rms_nfe_2": 2.5132674276828766, "loss_nfe_3": 0.3547529995441437, "edge_loss_nfe_3": 0.17915709614753722, "diag_loss_nfe_3": 0.17559590488672255, "jvp_rms_nfe_3": 2.0510569334030153, "loss_nfe_4": 0.3877130031585693, "edge_loss_nfe_4": 0.19501143097877502, "diag_loss_nfe_4": 0.19270157516002656, "jvp_rms_nfe_4": 2.470135527849197, "loss_nfe_5": 0.3525314390659332, "edge_loss_nfe_5": 0.17706592082977296, "diag_loss_nfe_5": 0.17546551525592805, "jvp_rms_nfe_5": 1.606951016187668, "raw_loss": 0.3729947149753571, "raw_edge_loss": 0.19341899901628495, "raw_diag_loss": 0.17957571774721146, "raw_jvp_rms": 2.3479965436458587, "raw_loss_nfe_1": 0.4013635516166687, "raw_edge_loss_nfe_1": 0.22706393301486968, "raw_diag_loss_nfe_1": 0.17429961562156676, "raw_jvp_rms_nfe_1": 2.7829211711883546, "raw_loss_nfe_2": 0.3686074197292328, "raw_edge_loss_nfe_2": 0.18846496045589448, "raw_diag_loss_nfe_2": 0.18014245927333833, "raw_jvp_rms_nfe_2": 2.6206358075141907, "raw_loss_nfe_3": 0.35482444763183596, "raw_edge_loss_nfe_3": 0.17929541170597077, "raw_diag_loss_nfe_3": 0.17552904039621353, "raw_jvp_rms_nfe_3": 2.1178201019763945, "raw_loss_nfe_4": 0.3876810848712921, "raw_edge_loss_nfe_4": 0.19514916241168975, "raw_diag_loss_nfe_4": 0.19253192245960235, "raw_jvp_rms_nfe_4": 2.551815760135651, "raw_loss_nfe_5": 0.35249707102775574, "raw_edge_loss_nfe_5": 0.17712152749300003, "raw_diag_loss_nfe_5": 0.1753755509853363, "raw_jvp_rms_nfe_5": 1.6667898774147034}
+{"step": 50000, "weights_source": "ema", "loss": 0.3709731566905975, "edge_loss": 0.19136443853378296, "diag_loss": 0.1796087145805359, "jvp_rms": 2.2817168676853177, "examples": 5120, "loss_nfe_1": 0.39239816665649413, "edge_loss_nfe_1": 0.21810964345932007, "diag_loss_nfe_1": 0.1742885172367096, "jvp_rms_nfe_1": 2.700194501876831, "loss_nfe_2": 0.3679115891456604, "edge_loss_nfe_2": 0.18770996630191802, "diag_loss_nfe_2": 0.18020162284374236, "jvp_rms_nfe_2": 2.536956226825714, "loss_nfe_3": 0.35463637113571167, "edge_loss_nfe_3": 0.1791027843952179, "diag_loss_nfe_3": 0.17553358674049377, "jvp_rms_nfe_3": 2.0643750607967375, "loss_nfe_4": 0.38753209710121156, "edge_loss_nfe_4": 0.19491184055805205, "diag_loss_nfe_4": 0.19262024462223054, "jvp_rms_nfe_4": 2.4871247947216033, "loss_nfe_5": 0.3523875594139099, "edge_loss_nfe_5": 0.17698795795440675, "diag_loss_nfe_5": 0.17539960145950317, "jvp_rms_nfe_5": 1.6199337542057037, "raw_loss": 0.37272677540779114, "raw_edge_loss": 0.19294691383838652, "raw_diag_loss": 0.17977986514568328, "raw_jvp_rms": 2.4338310968875883, "raw_loss_nfe_1": 0.39853897094726565, "raw_edge_loss_nfe_1": 0.2240687519311905, "raw_diag_loss_nfe_1": 0.1744702249765396, "raw_jvp_rms_nfe_1": 2.9251581192016602, "raw_loss_nfe_2": 0.3692192077636719, "raw_edge_loss_nfe_2": 0.18882804214954377, "raw_diag_loss_nfe_2": 0.18039117157459258, "raw_jvp_rms_nfe_2": 2.7281934142112734, "raw_loss_nfe_3": 0.35528815984725953, "raw_edge_loss_nfe_3": 0.1796249717473984, "raw_diag_loss_nfe_3": 0.17566319108009337, "raw_jvp_rms_nfe_3": 2.1835286557674407, "raw_loss_nfe_4": 0.3877676367759705, "raw_edge_loss_nfe_4": 0.19502896666526795, "raw_diag_loss_nfe_4": 0.1927386701107025, "raw_jvp_rms_nfe_4": 2.6403229236602783, "raw_loss_nfe_5": 0.3528199017047882, "raw_edge_loss_nfe_5": 0.1771838366985321, "raw_diag_loss_nfe_5": 0.17563606798648834, "raw_jvp_rms_nfe_5": 1.69195237159729}