Add ComposedRouter for multi-axis MoE gating
Route on several independent axes at once (e.g. energy x pdg), each with
its own expert count and hyperparameters. The joint gate is the outer
product of per-axis softmax gates, so it stays a partition of unity and
top1/balance_loss factor per-axis. Config uses flat axis{i}_{field} keys
in model.router (TOML/CLI friendly), also settable via repeatable
--router-axis flags.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
+50
-1
@@ -63,6 +63,43 @@ def _batch_size_estimate_dims(model_cfg: dict) -> tuple[int, int]:
|
||||
return model_cfg["hidden_dim"], model_cfg["n_blocks"]
|
||||
|
||||
|
||||
def _coerce_scalar(value: str) -> object:
|
||||
"""Best-effort str -> bool/int/float, else leave as str.
|
||||
|
||||
CLI flag values always arrive as strings; router kwargs like
|
||||
`n_experts` (int) or `temperature` (float) need to come out typed the
|
||||
same way a TOML file's native types would, since they're merged into
|
||||
the same `model.router` dict as file-sourced config.
|
||||
"""
|
||||
if value.lower() in ("true", "false"):
|
||||
return value.lower() == "true"
|
||||
try:
|
||||
return int(value)
|
||||
except ValueError:
|
||||
pass
|
||||
try:
|
||||
return float(value)
|
||||
except ValueError:
|
||||
pass
|
||||
return value
|
||||
|
||||
|
||||
def _parse_router_axis_flags(specs: list[str]) -> dict[str, object]:
|
||||
"""Parse repeated `--router-axis "type:key=val,key=val"` flags into
|
||||
`axis{i}_{field}` flat keys (see `_parse_composed_axes` in
|
||||
giant.model.network), indexed by flag order — the Nth `--router-axis`
|
||||
becomes axis N.
|
||||
"""
|
||||
out: dict[str, object] = {}
|
||||
for i, spec in enumerate(specs):
|
||||
axis_type, _, rest = spec.partition(":")
|
||||
out[f"axis{i}_type"] = axis_type
|
||||
for pair in filter(None, rest.split(",")):
|
||||
key, _, val = pair.partition("=")
|
||||
out[f"axis{i}_{key}"] = _coerce_scalar(val)
|
||||
return out
|
||||
|
||||
|
||||
_CEPH_PREDICTIONS = Path("/ceph/lbogner/geant_steps/predictions")
|
||||
|
||||
|
||||
@@ -175,6 +212,16 @@ def train(
|
||||
n_experts: Annotated[
|
||||
Optional[int], typer.Option("--n-experts", help="Number of routed experts")
|
||||
] = None,
|
||||
router_axis: Annotated[
|
||||
Optional[list[str]],
|
||||
typer.Option(
|
||||
"--router-axis",
|
||||
help="Composed-router axis spec 'type:key=val,key=val' (repeatable; "
|
||||
"Nth flag = axis N). Use with --router-type composed instead of "
|
||||
"--n-experts, e.g. --router-axis 'energy:n_experts=4' "
|
||||
"--router-axis 'pdg:n_experts=3,emb_dim=8'",
|
||||
),
|
||||
] = None,
|
||||
val_fraction: Annotated[
|
||||
Optional[float], typer.Option("--val-fraction", "-f")
|
||||
] = None,
|
||||
@@ -264,7 +311,7 @@ def train(
|
||||
}.items()
|
||||
if v is not None
|
||||
}
|
||||
cli_router = {
|
||||
cli_router: dict[str, object] = {
|
||||
k: v
|
||||
for k, v in {
|
||||
"enabled": router,
|
||||
@@ -273,6 +320,8 @@ def train(
|
||||
}.items()
|
||||
if v is not None
|
||||
}
|
||||
if router_axis:
|
||||
cli_router.update(_parse_router_axis_flags(router_axis))
|
||||
if cli_router:
|
||||
cli_model["router"] = cli_router
|
||||
cfg = gconfig.merge_cli_overrides(
|
||||
|
||||
@@ -43,6 +43,13 @@ DEFAULT_CONFIG: dict = {
|
||||
# (0.0 still trains a working router — the gate gets gradient
|
||||
# through the downstream flow loss like EnergyRouter's centers —
|
||||
# but only lambda_proc > 0 grounds it in the true `process` label)
|
||||
# type = "composed" routes on multiple axes at once (e.g. energy x
|
||||
# pdg), each with its own expert count/hyperparameters. Axes are
|
||||
# NOT in these defaults (there's no meaningful default axis list)
|
||||
# — set them as flat axis{i}_{field} keys instead of "n_experts",
|
||||
# e.g. axis0_type = "energy", axis0_n_experts = 4, axis1_type =
|
||||
# "pdg", axis1_n_experts = 3, axis1_emb_dim = 8. See
|
||||
# giant.model.network._parse_composed_axes / `--router-axis`.
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
+128
-19
@@ -1,5 +1,6 @@
|
||||
import inspect
|
||||
import math
|
||||
import re
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
@@ -444,6 +445,79 @@ class ProcessRouter(Router):
|
||||
return F.cross_entropy(self.logits(cond_cont, cond_cat), labels)
|
||||
|
||||
|
||||
class ComposedRouter(Router):
|
||||
"""Joint router over independent axes (e.g. energy x pdg), outer-product gated.
|
||||
|
||||
Wraps N already-built sub-routers, each free to have its own
|
||||
`n_experts` and hyperparameters (an `EnergyRouter(n_experts=4, ...)`
|
||||
composed with a `PdgRouter(n_experts=3, ...)` needs no axis to match
|
||||
the other's expert count). The joint gate is the outer product of the
|
||||
per-axis softmax gates, flattened to `(B, prod(n_experts_i))` — still a
|
||||
partition of unity, since each factor is one. Because the axes are
|
||||
routed independently, the joint argmax factors into the per-axis
|
||||
argmaxes, so `top1` (inherited from `Router`) costs no more than
|
||||
routing each axis alone despite the multiplicative expert count; the
|
||||
same is true of `balance_loss` (inherited, computed on the flattened
|
||||
joint gate — now one importance term per *joint* expert cell).
|
||||
|
||||
Not registered in `ROUTER_REGISTRY` / buildable via `build_router`,
|
||||
since those assume one `n_experts` int shared by a single router type;
|
||||
use `build_composed_router` instead, which resolves a list of per-axis
|
||||
specs (each independently typed and sized) through `build_router`.
|
||||
"""
|
||||
|
||||
def __init__(self, routers: list[Router]) -> None:
|
||||
if not routers:
|
||||
raise ValueError("ComposedRouter needs at least one sub-router")
|
||||
n_experts = 1
|
||||
for r in routers:
|
||||
n_experts *= r.n_experts
|
||||
super().__init__(n_experts)
|
||||
self.routers = nn.ModuleList(routers)
|
||||
|
||||
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
||||
joint = self.routers[0].gate(cond_cont, cond_cat) # (B, n_0)
|
||||
for router in self.routers[1:]:
|
||||
g = router.gate(cond_cont, cond_cat) # (B, n_i)
|
||||
joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(
|
||||
1
|
||||
) # (B, prod so far)
|
||||
return joint
|
||||
|
||||
def classify_loss(
|
||||
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
"""Sum of each sub-router's own classify_loss (0 for unsupervised axes)."""
|
||||
total = torch.zeros((), device=cond_cont.device)
|
||||
for router in self.routers:
|
||||
total = total + router.classify_loss(cond_cont, cond_cat, labels)
|
||||
return total
|
||||
|
||||
|
||||
def build_composed_router(specs: list[dict], **shared_kwargs) -> ComposedRouter:
|
||||
"""Build a `ComposedRouter` from a list of per-axis router specs.
|
||||
|
||||
Each spec is a `{"type": ..., "n_experts": ..., ...per-axis kwargs}`
|
||||
dict resolved through `build_router` exactly like a single-axis router
|
||||
config, so axes can differ in both expert count and hyperparameters
|
||||
(e.g. an energy axis's `temperature` vs a pdg axis's `emb_dim`).
|
||||
`shared_kwargs` (`pdg_vocab`, `mat_vocab`, ...) are merged under each
|
||||
spec, with the spec's own keys taking precedence.
|
||||
"""
|
||||
routers = [
|
||||
build_router(
|
||||
spec["type"],
|
||||
spec["n_experts"],
|
||||
**{
|
||||
**shared_kwargs,
|
||||
**{k: v for k, v in spec.items() if k not in ("type", "n_experts")},
|
||||
},
|
||||
)
|
||||
for spec in specs
|
||||
]
|
||||
return ComposedRouter(routers)
|
||||
|
||||
|
||||
class ExpertTrunk(nn.Module):
|
||||
"""One small expert: `input_proj -> ResBlock stack -> out_proj`.
|
||||
|
||||
@@ -669,6 +743,54 @@ _SEC_DECODER_MODEL_KEYS = {
|
||||
}
|
||||
|
||||
|
||||
_AXIS_KEY_RE = re.compile(r"^axis(\d+)_(.+)$")
|
||||
|
||||
|
||||
def _parse_composed_axes(router_cfg: dict) -> list[dict]:
|
||||
"""Regroup `axis{i}_{field}` flat keys into a list of per-axis spec dicts.
|
||||
|
||||
Flat keys (rather than a nested list-of-dicts) keep composed-router
|
||||
config expressible in the same one-level-of-nesting TOML/CLI shape as
|
||||
every other router option (`model.router` stays a flat table of
|
||||
scalars) — e.g. `axis0_type = "energy"`, `axis0_n_experts = 4`,
|
||||
`axis1_type = "pdg"`, `axis1_n_experts = 3`, `axis1_emb_dim = 8`.
|
||||
Axis indices must be contiguous from 0; order follows the index, not
|
||||
dict insertion order (TOML/CLI merging doesn't preserve it reliably).
|
||||
"""
|
||||
axes: dict[int, dict] = {}
|
||||
for key, value in router_cfg.items():
|
||||
m = _AXIS_KEY_RE.match(key)
|
||||
if m is None:
|
||||
continue
|
||||
idx, field = int(m.group(1)), m.group(2)
|
||||
axes.setdefault(idx, {})[field] = value
|
||||
missing = set(range(len(axes))) - axes.keys()
|
||||
if missing:
|
||||
raise ValueError(f"composed router config has gaps at axis indices {missing}")
|
||||
return [axes[i] for i in range(len(axes))]
|
||||
|
||||
|
||||
def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int) -> Router:
|
||||
"""Resolve one `model.router` config into a `Router`, single-axis or composed.
|
||||
|
||||
`router_cfg["type"] == "composed"` reads `axis{i}_{field}` flat keys
|
||||
(see `_parse_composed_axes`) instead of a single `type`/`n_experts` pair.
|
||||
"""
|
||||
shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
|
||||
if router_cfg["type"] == "composed":
|
||||
return build_composed_router(_parse_composed_axes(router_cfg), **shared_vocab)
|
||||
router_kwargs = {
|
||||
k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")
|
||||
}
|
||||
# Not every router needs these (EnergyRouter doesn't declare them, so
|
||||
# build_router's kwarg filtering drops them silently) but ProcessRouter
|
||||
# needs its own pdg/material embeddings sized to match the checkpoint's
|
||||
# vocab, same as the trunk's ConditionEncoder.
|
||||
router_kwargs.setdefault("pdg_vocab", pdg_vocab)
|
||||
router_kwargs.setdefault("mat_vocab", mat_vocab)
|
||||
return build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
|
||||
|
||||
|
||||
def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
|
||||
"""Construct (stage1, sec_decoder) from a persisted/CLI model_config dict.
|
||||
|
||||
@@ -679,36 +801,23 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
|
||||
"""
|
||||
router_cfg = model_config.get("router")
|
||||
if router_cfg and router_cfg.get("enabled"):
|
||||
router_kwargs = {
|
||||
k: v
|
||||
for k, v in router_cfg.items()
|
||||
if k not in ("enabled", "type", "n_experts")
|
||||
}
|
||||
# Not every router needs these (EnergyRouter doesn't declare them, so
|
||||
# build_router's kwarg filtering drops them silently) but
|
||||
# ProcessRouter needs its own pdg/material embeddings sized to match
|
||||
# the checkpoint's vocab, same as the trunk's ConditionEncoder.
|
||||
router_kwargs.setdefault("pdg_vocab", model_config["pdg_vocab"])
|
||||
router_kwargs.setdefault("mat_vocab", model_config["mat_vocab"])
|
||||
pdg_vocab = model_config["pdg_vocab"]
|
||||
mat_vocab = model_config["mat_vocab"]
|
||||
shared = dict(
|
||||
pdg_vocab=model_config["pdg_vocab"],
|
||||
mat_vocab=model_config["mat_vocab"],
|
||||
pdg_vocab=pdg_vocab,
|
||||
mat_vocab=mat_vocab,
|
||||
expert_hidden_dim=model_config.get("expert_hidden_dim", 128),
|
||||
expert_n_blocks=model_config.get("expert_n_blocks", 3),
|
||||
emb_dim=model_config.get("emb_dim", EMB_DIM),
|
||||
dropout=model_config.get("dropout", 0.1),
|
||||
)
|
||||
stage1 = RoutedDenoisingMLP(
|
||||
router=build_router(
|
||||
router_cfg["type"], router_cfg["n_experts"], **router_kwargs
|
||||
),
|
||||
router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab),
|
||||
k_max=model_config.get("k_max", K_MAX),
|
||||
**shared,
|
||||
)
|
||||
sec_decoder = RoutedSecondaryDecoder(
|
||||
router=build_router(
|
||||
router_cfg["type"], router_cfg["n_experts"], **router_kwargs
|
||||
),
|
||||
router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab),
|
||||
**shared,
|
||||
)
|
||||
return stage1, sec_decoder
|
||||
|
||||
@@ -4,6 +4,7 @@ import torch
|
||||
|
||||
from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
|
||||
from giant.model.network import (
|
||||
ComposedRouter,
|
||||
DenoisingMLP,
|
||||
EnergyRouter,
|
||||
PdgRouter,
|
||||
@@ -12,6 +13,7 @@ from giant.model.network import (
|
||||
RoutedDenoisingMLP,
|
||||
RoutedSecondaryDecoder,
|
||||
SecondaryDecoder,
|
||||
build_composed_router,
|
||||
build_models,
|
||||
build_router,
|
||||
)
|
||||
@@ -289,6 +291,224 @@ def test_build_models_routed_with_process_router():
|
||||
assert stage1.router.mat_emb.num_embeddings == 2
|
||||
|
||||
|
||||
# ── ComposedRouter ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_composed_router_n_experts_is_product():
|
||||
router = ComposedRouter(
|
||||
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
|
||||
)
|
||||
assert router.n_experts == 12
|
||||
|
||||
|
||||
def test_composed_router_gate_partition_of_unity():
|
||||
router = ComposedRouter(
|
||||
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
|
||||
)
|
||||
cond_cont, cond_cat = _cond(16, pdg=5)
|
||||
g = router.gate(cond_cont, cond_cat)
|
||||
assert g.shape == (16, 12)
|
||||
torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
|
||||
|
||||
|
||||
def test_composed_router_gate_is_outer_product_of_sub_gates():
|
||||
energy_router = EnergyRouter(n_experts=4)
|
||||
pdg_router = PdgRouter(n_experts=3, pdg_vocab=5)
|
||||
router = ComposedRouter([energy_router, pdg_router])
|
||||
cond_cont, cond_cat = _cond(16, pdg=5)
|
||||
|
||||
g_energy = energy_router.gate(cond_cont, cond_cat) # (16, 4)
|
||||
g_pdg = pdg_router.gate(cond_cont, cond_cat) # (16, 3)
|
||||
expected = (g_energy.unsqueeze(-1) * g_pdg.unsqueeze(1)).flatten(1) # (16, 12)
|
||||
|
||||
torch.testing.assert_close(router.gate(cond_cont, cond_cat), expected)
|
||||
|
||||
|
||||
def test_composed_router_top1_factors_into_per_axis_argmax():
|
||||
"""Joint argmax over the outer product must equal the pair of per-axis
|
||||
argmaxes, flattened with the same row-major index convention as gate()."""
|
||||
energy_router = EnergyRouter(n_experts=4)
|
||||
pdg_router = PdgRouter(n_experts=3, pdg_vocab=5)
|
||||
router = ComposedRouter([energy_router, pdg_router])
|
||||
cond_cont, cond_cat = _cond(16, pdg=5)
|
||||
|
||||
joint_idx = router.top1(cond_cont, cond_cat)
|
||||
energy_idx = energy_router.top1(cond_cont, cond_cat)
|
||||
pdg_idx = pdg_router.top1(cond_cont, cond_cat)
|
||||
expected = energy_idx * pdg_router.n_experts + pdg_idx
|
||||
|
||||
assert torch.equal(joint_idx, expected)
|
||||
|
||||
|
||||
def test_composed_router_supports_different_expert_counts_per_axis():
|
||||
router = ComposedRouter(
|
||||
[EnergyRouter(n_experts=5), PdgRouter(n_experts=2, pdg_vocab=5)]
|
||||
)
|
||||
assert router.n_experts == 10
|
||||
cond_cont, cond_cat = _cond(8, pdg=5)
|
||||
assert router.gate(cond_cont, cond_cat).shape == (8, 10)
|
||||
|
||||
|
||||
def test_composed_router_classify_loss_sums_sub_router_losses():
|
||||
"""energy/pdg both default to zero, so the composed loss should too."""
|
||||
router = ComposedRouter(
|
||||
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
|
||||
)
|
||||
cond_cont, cond_cat = _cond(16, pdg=5)
|
||||
labels = torch.randint(0, 4, (16,))
|
||||
loss = router.classify_loss(cond_cont, cond_cat, labels)
|
||||
assert loss.shape == ()
|
||||
assert loss.item() == 0.0
|
||||
|
||||
|
||||
def test_composed_router_rejects_empty_router_list():
|
||||
try:
|
||||
ComposedRouter([])
|
||||
except ValueError:
|
||||
return
|
||||
raise AssertionError("expected ValueError for empty router list")
|
||||
|
||||
|
||||
def test_composed_router_not_in_registry():
|
||||
assert "composed" not in ROUTER_REGISTRY
|
||||
|
||||
|
||||
# ── _parse_composed_axes (axis{i}_{field} flat-key config convention) ───────
|
||||
|
||||
|
||||
def test_parse_composed_axes_groups_indexed_keys():
|
||||
from giant.model.network import _parse_composed_axes
|
||||
|
||||
router_cfg = {
|
||||
"enabled": True,
|
||||
"type": "composed",
|
||||
"axis0_type": "energy",
|
||||
"axis0_n_experts": 4,
|
||||
"axis0_temperature": 0.3,
|
||||
"axis1_type": "pdg",
|
||||
"axis1_n_experts": 3,
|
||||
"axis1_emb_dim": 6,
|
||||
}
|
||||
axes = _parse_composed_axes(router_cfg)
|
||||
assert axes == [
|
||||
{"type": "energy", "n_experts": 4, "temperature": 0.3},
|
||||
{"type": "pdg", "n_experts": 3, "emb_dim": 6},
|
||||
]
|
||||
|
||||
|
||||
def test_parse_composed_axes_ignores_unrelated_keys():
|
||||
from giant.model.network import _parse_composed_axes
|
||||
|
||||
router_cfg = {
|
||||
"enabled": True,
|
||||
"type": "composed",
|
||||
"lambda_balance": 0.0,
|
||||
"axis0_type": "energy",
|
||||
"axis0_n_experts": 4,
|
||||
}
|
||||
axes = _parse_composed_axes(router_cfg)
|
||||
assert axes == [{"type": "energy", "n_experts": 4}]
|
||||
|
||||
|
||||
def test_parse_composed_axes_raises_on_index_gap():
|
||||
from giant.model.network import _parse_composed_axes
|
||||
|
||||
router_cfg = {
|
||||
"type": "composed",
|
||||
"axis0_type": "energy",
|
||||
"axis0_n_experts": 4,
|
||||
# axis1 missing entirely
|
||||
"axis2_type": "pdg",
|
||||
"axis2_n_experts": 3,
|
||||
}
|
||||
try:
|
||||
_parse_composed_axes(router_cfg)
|
||||
except ValueError:
|
||||
return
|
||||
raise AssertionError("expected ValueError for a gap in axis indices")
|
||||
|
||||
|
||||
def test_build_composed_router_resolves_per_axis_specs():
|
||||
router = build_composed_router(
|
||||
[
|
||||
{"type": "energy", "n_experts": 4, "temperature": 0.3},
|
||||
{"type": "pdg", "n_experts": 3, "emb_dim": 6},
|
||||
],
|
||||
pdg_vocab=5,
|
||||
mat_vocab=2,
|
||||
)
|
||||
assert isinstance(router, ComposedRouter)
|
||||
assert router.n_experts == 12
|
||||
energy_router, pdg_router = router.routers
|
||||
assert isinstance(energy_router, EnergyRouter)
|
||||
assert energy_router.temperature == 0.3
|
||||
assert isinstance(pdg_router, PdgRouter)
|
||||
assert pdg_router.pdg_emb.num_embeddings == 5
|
||||
assert pdg_router.pdg_emb.embedding_dim == 6
|
||||
|
||||
|
||||
def test_build_models_routed_with_composed_router():
|
||||
model_config = dict(
|
||||
pdg_vocab=5,
|
||||
mat_vocab=2,
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
expert_hidden_dim=16,
|
||||
expert_n_blocks=2,
|
||||
router={
|
||||
"enabled": True,
|
||||
"type": "composed",
|
||||
"axis0_type": "energy",
|
||||
"axis0_n_experts": 4,
|
||||
"axis1_type": "pdg",
|
||||
"axis1_n_experts": 3,
|
||||
},
|
||||
)
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
assert isinstance(stage1, RoutedDenoisingMLP)
|
||||
assert isinstance(stage1.router, ComposedRouter)
|
||||
assert len(stage1.experts) == 12
|
||||
assert len(sec_decoder.experts) == 12
|
||||
# stage1 and sec_decoder must not share router weights (same convention
|
||||
# as the single-axis routers built by build_models).
|
||||
assert stage1.router is not sec_decoder.router
|
||||
|
||||
|
||||
def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow():
|
||||
from giant.sample import sample_flow, sample_secondaries
|
||||
|
||||
model_config = dict(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
expert_hidden_dim=8,
|
||||
expert_n_blocks=1,
|
||||
router={
|
||||
"enabled": True,
|
||||
"type": "composed",
|
||||
"axis0_type": "energy",
|
||||
"axis0_n_experts": 2,
|
||||
"axis1_type": "pdg",
|
||||
"axis1_n_experts": 2,
|
||||
},
|
||||
)
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
B = 5
|
||||
cond_cont, cond_cat = _cond(B, pdg=3, mat=2)
|
||||
stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2)
|
||||
assert stage1_norm.shape == (B, X_DIM)
|
||||
assert n_sec_pred.shape == (B,)
|
||||
|
||||
sec_cont, sec_type_emb, sec_valid = sample_secondaries(
|
||||
sec_decoder, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
|
||||
)
|
||||
assert sec_cont.shape == (B, K_MAX, 4)
|
||||
assert sec_valid.shape == (B, K_MAX)
|
||||
|
||||
|
||||
# ── RoutedDenoisingMLP ───────────────────────────────────────────────────────
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user