- 移除 JEPA/lejepa-identifiability 子模块 gitlink - 移除 research/multiply/MultiPLY 子模块 gitlink - 删除 .gitmodules(不再有外部 URL 依赖) - 两个目录内容作为普通文件纳入主仓库追踪 - 删除各自内部 .git 目录,消除嵌套 git 仓库
This commit is contained in:
@@ -0,0 +1,56 @@
|
||||
"""Nonlinear mixing functions."""
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
# ── 2D mixing functions ──────────────────────────────────────────────────────
|
||||
|
||||
def mix_spiral(z):
|
||||
"""g(z) = R(π‖z‖) z — measure-preserving spiral diffeomorphism."""
|
||||
norms = z.norm(dim=-1) * torch.pi
|
||||
c, s = norms.cos(), norms.sin()
|
||||
R = torch.stack([torch.stack([c, -s], dim=-1),
|
||||
torch.stack([s, c], dim=-1)], dim=-2)
|
||||
return (R @ z.unsqueeze(-1)).squeeze(-1)
|
||||
|
||||
|
||||
def mix_banana(z):
|
||||
"""Banana: x0 = z0, x1 = z1 + z0²."""
|
||||
return torch.stack([z[..., 0], z[..., 1] + z[..., 0] ** 2], dim=-1)
|
||||
|
||||
|
||||
def mix_sinusoid(z):
|
||||
"""Sinusoidal shear: x0 = z0 + sin(1.5 z1), x1 = z1."""
|
||||
return torch.stack([z[..., 0] + torch.sin(1.5 * z[..., 1]), z[..., 1]], dim=-1)
|
||||
|
||||
|
||||
MIXINGS_2D = {
|
||||
"spiral": mix_spiral,
|
||||
"banana": mix_banana,
|
||||
"sinusoid": mix_sinusoid,
|
||||
# "nvp" handled via make_coupling_mixing(N=2, n_layers=...)
|
||||
}
|
||||
|
||||
|
||||
# ── Coupling-layer mixing (any dimension) ────────────────────────────────────
|
||||
|
||||
def make_coupling_mixing(N, n_layers=4, seed=1337, device="cuda"):
|
||||
"""RealNVP-style coupling layers. Works for any even N (including N=2)."""
|
||||
half = N // 2
|
||||
torch.manual_seed(seed)
|
||||
Ws = []
|
||||
for _ in range(n_layers):
|
||||
W, _ = torch.linalg.qr(torch.randn(half, half, device=device))
|
||||
Ws.append(W * 2.0)
|
||||
|
||||
def mix(z):
|
||||
for i, W in enumerate(Ws):
|
||||
z1, z2 = z[..., :half], z[..., half:]
|
||||
if i % 2 == 0:
|
||||
z2 = z2 + torch.tanh(z1 @ W)
|
||||
else:
|
||||
z1 = z1 + torch.tanh(z2 @ W)
|
||||
z = torch.cat([z1, z2], dim=-1)
|
||||
return z
|
||||
|
||||
return mix
|
||||
Reference in New Issue
Block a user