feat: 目标写法开关 mf_target(mf/lin/sq/mid)，默认 mf 且原地路径逐位不变；float32 接线用解析高斯场自证

tests/test_meanflow_targets.py

@@ -9,6 +9,7 @@
 import unittest
 
 import torch
+from torch.func import jvp
 
 LAM = 0.5
 
@@ -108,5 +109,64 @@ class MeanFlowTargetFixedPointTest(unittest.TestCase):
         self.assertLess(half['mid'], full['mid'])
 
 
+class AvgVelocityTargetWiringTest(unittest.TestCase):
+    """把三种写法接进 ml/model/meanflow.py 之后，再用同一个解析场验一遍接线：
+    浮点从 float64 换到 float32，量级判据放宽，但"mf 无偏 / sq 只线性收敛"必须照样成立，
+    否则说明 z_c、切向量、r 槽换位这几处接错了。"""
+    SIG = 2.0
+
+    class Stub:
+        def __init__(self, sig2):
+            self.sig2 = sig2
+
+        def s2(self, t):
+            return (1 - t) ** 2 * self.sig2 + t ** 2
+
+        def decode(self, mem, z, r, t):
+            return z * (1 - torch.sqrt(self.s2(r) / self.s2(t)))[:, None, None] / (t - r)[:, None, None]
+
+    def bias(self, target, r, t, lam=0.5):
+        from ml.model.meanflow import avg_velocity_target
+        f = Field(self.SIG)
+        z = torch.tensor([[0.7] * 4, [0.7] * 4], dtype=torch.float32)          # [2,4,6]? 形状自定，逐元素即可
+        z = torch.full((2, 4, 6), 0.7)
+        rt = torch.tensor([r, r]); tt = torch.tensor([t, t])
+        xh = (z * (1 - tt)[:, None, None] * self.SIG ** 2 / f.s2(tt)[:, None, None])
+        eh = (z * tt[:, None, None] / f.s2(tt)[:, None, None])
+        u, tgt, _ = avg_velocity_target(self.Stub(self.SIG ** 2), None, xh, eh, rt, tt, target, lam)
+        truth = f.u(torch.tensor(0.7, dtype=torch.float32), torch.tensor(r), torch.tensor(t))
+        return float(((tgt - truth).abs() / truth.abs()).mean()), float(truth)
+
+    def test_wiring_keeps_mf_unbiased_and_sq_first_order(self):
+        b_mf, _ = self.bias('mf', 0.8, 1.0)
+        self.assertLess(b_mf, 1e-4)
+        rel = []
+        for s in (0.2, 0.1, 0.05):
+            rel.append({k: self.bias(k, 1.0 - s, 1.0)[0] for k in ('sq', 'lin', 'mid', 'mf')})
+        print('float32 接线后的偏差：', [{k: round(v, 4) for k, v in d.items()} for d in rel])
+        r_sq = rel[1]['sq'] / rel[0]['sq']
+        r_mid = rel[1]['mid'] / rel[0]['mid']
+        self.assertGreater(r_sq, 0.3)          # 只线性收敛 ⇒ 结构性偏差
+        self.assertLess(r_mid, 0.35)           # 二次收敛 ⇒ 只剩积分误差
+        self.assertLess(rel[0]['lin'], rel[0]['sq'])
+
+    def test_default_path_is_bitwise_identical(self):
+        """target='mf' 必须与重构前的原地写法逐位相同（默认训练路径不许被这次改动碰到）"""
+        from ml.model.meanflow import avg_velocity_target
+        torch.manual_seed(0)
+        m = self.Stub(2.0)
+        x = torch.randn(3, 5, 6); e = torch.randn(3, 5, 6)
+        r = torch.rand(3) * 0.4; t = 0.6 + torch.rand(3) * 0.4
+        u_new, tgt_new, z_new = avg_velocity_target(m, None, x, e, r, t, 'mf', 0.5)
+        fn = lambda a, b, c: m.decode(None, a, b, c)
+        z = (1 - t)[:, None, None] * x + t[:, None, None] * e
+        v = e - x
+        u0, dudt = jvp(fn, (z, r, t), (v, torch.zeros_like(r), torch.ones_like(t)))
+        tgt0 = (v - (t - r)[:, None, None] * dudt).detach()
+        for name, a_, b_ in (('u', u_new, u0), ('u_tgt', tgt_new, tgt0), ('z', z_new, z)):
+            with self.subTest(name=name):
+                torch.testing.assert_close(a_, b_, rtol=0, atol=0)
+
+
 if __name__ == '__main__':
     unittest.main()