From 0f95e0eaae0f2ab4211fb4e7e24558bbda0aa343 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Fri, 14 Aug 2026 09:41:50 +0200 Subject: [PATCH] Make ResBlock's conditioning-injection mechanism selectable (gitea #34) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ResBlock injected conditioning exactly one way — h = linear1(h) + cond_proj(cond), a conditional bias, the weakest standard option for a model whose entire job is to be conditional. Adds BLOCK_REGISTRY (giant/model/layers.py), mirroring the TRUNK_REGISTRY/ROUTER_REGISTRY registry+factory idiom (gitea #33), with two new drop-in alternatives: FilmResBlock (per-channel scale+shift modulating the norm output, zero-init so conditioning has no effect at construction) and AdaLNResBlock (DiT-style AdaLN-Zero — the norm's own affine is replaced by a conditioning-derived scale/shift, plus a zero-init gate on the residual branch, making the block the exact identity function at init). Selected per stage via a new stage{1,2}_model.trunk.block_conditioning config leaf ("add" | "film" | "adaln", default "add"), threaded through build_trunk/build_expert_body/RoutedTrunk and the three stage model constructors. Default stays "add" and ResBlock's body is unchanged, so existing configs/checkpoints are bit-identical to before this change. Decided during planning: the new field lives on the existing TrunkConfig rather than a new top-level block/blocks config section; the WGAN CriticModel (which builds its own ResBlock stack outside TRUNK_REGISTRY) and the issue's mentioned blocks.norm/blocks.activation axes are both left out of scope. Co-Authored-By: Claude Opus 5 --- giant/config.py | 14 +++++-- giant/model/builders.py | 3 ++ giant/model/layers.py | 83 ++++++++++++++++++++++++++++++++++++++ giant/model/models.py | 18 ++++++++- giant/model/network.py | 17 +++++++- giant/model/trunks.py | 41 +++++++++++++++---- tests/test_config.py | 29 ++++++++++++++ tests/test_router.py | 89 +++++++++++++++++++++++++++++++++++++++++ 8 files changed, 280 insertions(+), 14 deletions(-) diff --git a/giant/config.py b/giant/config.py index 71c83de..0bf1fc8 100644 --- a/giant/config.py +++ b/giant/config.py @@ -346,17 +346,23 @@ class TrunkConfig: stage, unaffected by this block. A future body's own hyperparameters (e.g. a transformer's `n_heads`/`n_layers`) would get their own sibling field here, matching how `flow`/`ddpm`/`wgan` already coexist selected by - `generator`.""" + `generator`. + + `block_conditioning` selects each body's conditioning-injection mechanism + from `giant.model.layers.BLOCK_REGISTRY` — `"add"` (default, today's + conditional-bias `ResBlock`, bit-identical to pre-gitea-#34 behaviour), + `"film"`, or `"adaln"`.""" type: str = "resmlp" + block_conditioning: str = "add" @classmethod def from_dict(cls, d: dict | None) -> "TrunkConfig": d = d or {} - return cls(type=d.get("type", "resmlp")) + return cls(type=d.get("type", "resmlp"), block_conditioning=d.get("block_conditioning", "add")) def to_dict(self) -> dict: - return {"type": self.type} + return {"type": self.type, "block_conditioning": self.block_conditioning} @dataclass(frozen=True) @@ -1431,6 +1437,8 @@ _OUT_DIR_NAME_CANDIDATES = [ ), ("stage1_trunk_type", _path_candidate("stage1_model.trunk.type", "s1t-")), ("stage2_trunk_type", _path_candidate("stage2_model.trunk.type", "s2t-")), + ("stage1_block_cond", _path_candidate("stage1_model.trunk.block_conditioning", "s1bc-")), + ("stage2_block_cond", _path_candidate("stage2_model.trunk.block_conditioning", "s2bc-")), ("stage1_router", _router_candidate("stage1_model", "s1")), ("stage2_router", _router_candidate("stage2_model", "s2")), ( diff --git a/giant/model/builders.py b/giant/model/builders.py index f5c689d..9fe2ca1 100644 --- a/giant/model/builders.py +++ b/giant/model/builders.py @@ -84,6 +84,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: noise_dim=s1_spec.wgan.noise_dim, router=stage1_router, trunk_type=s1_spec.trunk.type, + block_conditioning=s1_spec.trunk.block_conditioning, n_sec_head_k_max=n_sec_head_k_max, cond_enc=shared_cond_enc, ) @@ -123,6 +124,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: k_max=k_max, router=stage2_router, trunk_type=s2_spec.trunk.type, + block_conditioning=s2_spec.trunk.block_conditioning, build_n_sec_head=n_sec_owner != "stage1", particle_type_cfg=particle_type_cfg, history=ar_cfg.history, @@ -151,6 +153,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: k_max=k_max, router=stage2_router, trunk_type=s2_spec.trunk.type, + block_conditioning=s2_spec.trunk.block_conditioning, build_n_sec_head=n_sec_owner != "stage1", particle_type_cfg=particle_type_cfg, cond_enc=shared_cond_enc, diff --git a/giant/model/layers.py b/giant/model/layers.py index 84834b1..64591ea 100644 --- a/giant/model/layers.py +++ b/giant/model/layers.py @@ -57,6 +57,26 @@ class ContextAdapter(nn.Module): return torch.tanh(self.proj(x)) +BLOCK_REGISTRY: dict[str, type[nn.Module]] = {} + + +def register_block(name: str): + def decorator(cls: type[nn.Module]) -> type[nn.Module]: + BLOCK_REGISTRY[name] = cls + return cls + + return decorator + + +def build_block(name: str, dim: int, cond_dim: int, dropout: float = 0.0) -> nn.Module: + """Factory: look up a registered conditioning-injection block by name and + construct one instance — `trunk.block_conditioning` (gitea #34).""" + if name not in BLOCK_REGISTRY: + raise ValueError(f"unknown block conditioning type {name!r}; available: {sorted(BLOCK_REGISTRY)}") + return BLOCK_REGISTRY[name](dim, cond_dim, dropout) + + +@register_block("add") class ResBlock(nn.Module): def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None: super().__init__() @@ -74,3 +94,66 @@ class ResBlock(nn.Module): h = self.dropout(h) h = self.linear2(h) return x + h + + +@register_block("film") +class FilmResBlock(nn.Module): + """FiLM conditioning (Perez et al. 2018): a per-channel scale+shift + modulates the normalized features, on top of the norm's own affine — + an *additional* modulation, unlike `AdaLNResBlock` below, which replaces + the norm's affine outright. `film_proj` is zero-initialized so + `gamma=beta=0` at construction — conditioning has no effect on the + output until training moves it, a stable starting point (though not a + literal identity block, since `linear1`/`linear2` aren't zero-init).""" + + def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None: + super().__init__() + self.norm = nn.LayerNorm(dim) + self.linear1 = nn.Linear(dim, dim) + self.film_proj = nn.Linear(cond_dim, 2 * dim) + nn.init.zeros_(self.film_proj.weight) + nn.init.zeros_(self.film_proj.bias) + self.act = nn.SiLU() + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim, dim) + + def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: + h = self.norm(x) + gamma, beta = self.film_proj(cond).chunk(2, dim=-1) + h = h * (1 + gamma) + beta + h = self.linear1(h) + h = self.act(h) + h = self.dropout(h) + h = self.linear2(h) + return x + h + + +@register_block("adaln") +class AdaLNResBlock(nn.Module): + """AdaLN-Zero conditioning (DiT, Peebles & Xie 2022): the norm's own + affine is replaced by a conditioning-derived scale/shift, and the + residual branch is scaled by a conditioning-derived gate. `adaln_proj` + is zero-initialized, so `scale=shift=gate=0` at construction — the block + is the exact identity function at init (`x + 0 * h' == x`), regardless + of `x`/`cond`.""" + + def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None: + super().__init__() + self.norm = nn.LayerNorm(dim, elementwise_affine=False) + self.linear1 = nn.Linear(dim, dim) + self.adaln_proj = nn.Linear(cond_dim, 3 * dim) + nn.init.zeros_(self.adaln_proj.weight) + nn.init.zeros_(self.adaln_proj.bias) + self.act = nn.SiLU() + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim, dim) + + def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: + h = self.norm(x) + scale, shift, gate = self.adaln_proj(cond).chunk(3, dim=-1) + h = h * (1 + scale) + shift + h = self.linear1(h) + h = self.act(h) + h = self.dropout(h) + h = self.linear2(h) + return x + gate * h diff --git a/giant/model/models.py b/giant/model/models.py index ee3def7..9dcceb7 100644 --- a/giant/model/models.py +++ b/giant/model/models.py @@ -93,6 +93,7 @@ class Stage1Model(nn.Module): noise_dim: int = 64, router: Router | None = None, trunk_type: str = "resmlp", + block_conditioning: str = "add", n_sec_head_k_max: int | None = None, cond_enc: ConditionEncoder | None = None, ) -> None: @@ -108,7 +109,9 @@ class Stage1Model(nn.Module): self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim in_dim = noise_dim if generator == "wgan" else x_dim - self.trunk = build_trunk(router, trunk_type, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout) + self.trunk = build_trunk( + router, trunk_type, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout, block_conditioning + ) self.n_sec_head = None if n_sec_head_k_max is not None: self.n_sec_head = nn.Sequential( @@ -187,6 +190,7 @@ class Stage2OneShot(nn.Module): k_max: int = K_MAX, router: Router | None = None, trunk_type: str = "resmlp", + block_conditioning: str = "add", build_n_sec_head: bool = True, particle_type_cfg: dict | None = None, cond_enc: ConditionEncoder | None = None, @@ -214,7 +218,15 @@ class Stage2OneShot(nn.Module): merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim in_dim = noise_dim if generator == "wgan" else sec_dim self.trunk = build_trunk( - router, trunk_type, in_dim, sec_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout + router, + trunk_type, + in_dim, + sec_dim, + hidden_dim, + n_res_blocks, + merged_cond_dim, + dropout, + block_conditioning, ) self.n_sec_head = None if build_n_sec_head: @@ -333,6 +345,7 @@ class Stage2Autoregressive(nn.Module): k_max: int = K_MAX, router: Router | None = None, trunk_type: str = "resmlp", + block_conditioning: str = "add", build_n_sec_head: bool = True, particle_type_cfg: dict | None = None, history: str = "markov", @@ -392,6 +405,7 @@ class Stage2Autoregressive(nn.Module): n_res_blocks, merged_cond_dim, dropout, + block_conditioning, ) self.n_sec_head = None diff --git a/giant/model/network.py b/giant/model/network.py index 2cb2341..fc558ed 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -11,7 +11,17 @@ from giant.model._legacy import _migrate_legacy_model_config, migrate_legacy_sta from giant.model.builders import build_critics, build_models from giant.model.encoders import ConditionEncoder, cat_col_layout from giant.model.history import AttentionHistory, HistoryEncoder, MarkovHistory, _CausalAttnBlock -from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding, _make_axis_mlp +from giant.model.layers import ( + BLOCK_REGISTRY, + AdaLNResBlock, + ContextAdapter, + FilmResBlock, + ResBlock, + SinusoidalEmbedding, + _make_axis_mlp, + build_block, + register_block, +) from giant.model.models import ( CriticModel, Stage1Model, @@ -47,13 +57,16 @@ from giant.model.trunks import ( ) __all__ = [ + "AdaLNResBlock", "AttentionHistory", + "BLOCK_REGISTRY", "ComposedRouter", "ConditionEncoder", "ContextAdapter", "CriticModel", "EnergyRouter", "ExpertTrunk", + "FilmResBlock", "HistoryEncoder", "MarkovHistory", "PdgRouter", @@ -75,6 +88,7 @@ __all__ = [ "_migrate_legacy_model_config", "_parse_composed_axes", "_route_forward", + "build_block", "build_composed_router", "build_critics", "build_expert_body", @@ -83,6 +97,7 @@ __all__ = [ "build_trunk", "cat_col_layout", "migrate_legacy_state_dict", + "register_block", "register_router", "register_trunk", "resolve_type_n_classes", diff --git a/giant/model/trunks.py b/giant/model/trunks.py index ac31942..415d64d 100644 --- a/giant/model/trunks.py +++ b/giant/model/trunks.py @@ -12,7 +12,7 @@ free — no separate "routed transformer trunk" class needed. import torch import torch.nn as nn -from giant.model.layers import ResBlock +from giant.model.layers import build_block from giant.model.routers import Router TRUNK_REGISTRY: dict[str, type[nn.Module]] = {} @@ -34,14 +34,19 @@ def build_expert_body( n_blocks: int, cond_dim: int, dropout: float = 0.0, + block_conditioning: str = "add", ) -> nn.Module: """Factory: look up a registered trunk body by name and construct one instance of it — used both for a standalone (unrouted) trunk and for each - expert inside a `RoutedTrunk`.""" + expert inside a `RoutedTrunk`. `block_conditioning` selects the + `BLOCK_REGISTRY` entry each body's internal `ResBlock`-family blocks use + (`trunk.block_conditioning`, gitea #34) — an optional trailing kwarg a + future non-`ResBlock`-based body can simply ignore, same idiom as + `Trunk.forward`'s accept-and-ignore `cond_cont`/`cond_cat`.""" if name not in TRUNK_REGISTRY: raise ValueError(f"unknown trunk type {name!r}; available: {sorted(TRUNK_REGISTRY)}") cls = TRUNK_REGISTRY[name] - return cls(in_dim, out_dim, hidden_dim, n_blocks, cond_dim, dropout) + return cls(in_dim, out_dim, hidden_dim, n_blocks, cond_dim, dropout, block_conditioning=block_conditioning) @register_trunk("resmlp") @@ -65,11 +70,14 @@ class ExpertTrunk(nn.Module): n_blocks: int, cond_dim: int, dropout: float = 0.0, + block_conditioning: str = "add", ) -> None: super().__init__() self.out_dim = out_dim self.input_proj = nn.Linear(in_dim, hidden_dim) - self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_blocks)]) + self.blocks = nn.ModuleList( + [build_block(block_conditioning, hidden_dim, cond_dim, dropout) for _ in range(n_blocks)] + ) self.out_proj = nn.Linear(hidden_dim, out_dim) def forward( @@ -145,12 +153,22 @@ class RoutedTrunk(Trunk): n_res_blocks: int, cond_dim: int, dropout: float = 0.0, + block_conditioning: str = "add", ) -> None: super().__init__() self.router = router self.experts = nn.ModuleList( [ - build_expert_body(trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) + build_expert_body( + trunk_type, + in_dim, + out_dim, + hidden_dim, + n_res_blocks, + cond_dim, + dropout, + block_conditioning=block_conditioning, + ) for _ in range(router.n_experts) ] ) @@ -174,6 +192,7 @@ def build_trunk( n_res_blocks: int, cond_dim: int, dropout: float = 0.0, + block_conditioning: str = "add", ) -> nn.Module: """Build a stage's trunk: `trunk_type` (a `TRUNK_REGISTRY` key, e.g. `"resmlp"`) selects the expert body architecture; `router`, if given, @@ -182,8 +201,14 @@ def build_trunk( class), which is what makes an unrouted trunk's state-dict keys land directly under `trunk.*` instead of `trunk.experts.0.*` (see `giant.model._legacy.migrate_legacy_state_dict`, which assumes exactly - this flat layout for a v0.2 monolithic checkpoint). + this flat layout for a v0.2 monolithic checkpoint). `block_conditioning` + (a `BLOCK_REGISTRY` key, e.g. `"add"`/`"film"`/`"adaln"`) selects each + body's conditioning-injection mechanism (gitea #34). """ if router is not None: - return RoutedTrunk(router, trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) - return build_expert_body(trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) + return RoutedTrunk( + router, trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout, block_conditioning + ) + return build_expert_body( + trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout, block_conditioning + ) diff --git a/tests/test_config.py b/tests/test_config.py index e51a555..ea82ba5 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -85,6 +85,15 @@ def test_trunk_config_defaults_to_resmlp_for_both_stages(): assert gconfig.DEFAULT_CONFIG["stage2_model"]["trunk"]["type"] == "resmlp" +def test_trunk_config_defaults_block_conditioning_to_add_for_both_stages(): + """gitea #34: a pre-existing config with no `block_conditioning` key + must reproduce today's additive-bias behaviour exactly.""" + assert gconfig.Stage1ModelConfig().trunk.block_conditioning == "add" + assert gconfig.Stage2ModelConfig().trunk.block_conditioning == "add" + assert gconfig.DEFAULT_CONFIG["stage1_model"]["trunk"]["block_conditioning"] == "add" + assert gconfig.DEFAULT_CONFIG["stage2_model"]["trunk"]["block_conditioning"] == "add" + + def test_particle_type_config_n_classes_defaults_to_zero_and_round_trips(): """gitea #29: n_classes=0 means "inherit conditioning.particle.emb_dim" — the default must stay 0 so an existing config.toml with no @@ -936,6 +945,26 @@ def test_validate_config_keys_rejects_unknown_trunk_key(): assert "type" in str(e) +def test_validate_config_keys_allows_block_conditioning(): + cfg = _cfg_with( + **{ + "stage1_model.trunk.block_conditioning": "film", + "stage2_model.trunk.block_conditioning": "adaln", + } + ) + gconfig.validate_config_keys(cfg) # must not raise + + +def test_validate_config_keys_rejects_unknown_block_conditioning_key(): + cfg = _cfg_with(**{"stage1_model.trunk.block_conditioning_o": "film"}) # typo + try: + gconfig.validate_config_keys(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "stage1_model.trunk.block_conditioning_o" in str(e) + assert "block_conditioning" in str(e) + + def test_merge_cli_overrides_rejects_typo_in_toml_file(tmp_path): path = tmp_path / "config.toml" path.write_text("[meta]\nconfig_version = 3\n\n[stage1_model]\nn_res_block = 12\n") diff --git a/tests/test_router.py b/tests/test_router.py index 3c77c64..af1b403 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -5,16 +5,21 @@ import torch from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM from giant.model.network import ( + BLOCK_REGISTRY, TRUNK_REGISTRY, + AdaLNResBlock, ComposedRouter, EnergyRouter, ExpertTrunk, + FilmResBlock, PdgRouter, ProcessRouter, ROUTER_REGISTRY, + ResBlock, RoutedTrunk, Stage1Model, Stage2OneShot, + build_block, build_composed_router, build_expert_body, build_models, @@ -162,6 +167,52 @@ def test_build_expert_body_unknown_type_raises(): raise AssertionError("expected ValueError for unknown trunk type") +# ── BLOCK_REGISTRY / build_block (gitea #34) ──────────────────────────────── + + +def test_block_registry_has_add_film_adaln(): + assert BLOCK_REGISTRY["add"] is ResBlock + assert BLOCK_REGISTRY["film"] is FilmResBlock + assert BLOCK_REGISTRY["adaln"] is AdaLNResBlock + + +def test_build_block_unknown_type_raises(): + try: + build_block("nonexistent", dim=8, cond_dim=4) + except ValueError: + return + raise AssertionError("expected ValueError for unknown block conditioning type") + + +@pytest.mark.parametrize("block_type", ["add", "film", "adaln"]) +def test_block_forward_shape(block_type): + block = build_block(block_type, dim=8, cond_dim=4) + x = torch.randn(5, 8) + cond = torch.randn(5, 4) + out = block(x, cond) + assert out.shape == (5, 8) + + +def test_film_res_block_output_invariant_to_cond_at_init(): + """Zero-initialized film_proj means gamma=beta=0 at construction, so the + output must not depend on which cond is passed in.""" + block = FilmResBlock(dim=8, cond_dim=4) + x = torch.randn(5, 8) + cond_a = torch.randn(5, 4) + cond_b = torch.randn(5, 4) + torch.testing.assert_close(block(x, cond_a), block(x, cond_b)) + + +def test_adaln_res_block_is_identity_at_init(): + """Zero-initialized adaln_proj means scale=shift=gate=0 at construction, + so the block must be the exact identity function (the 'Zero' in + AdaLN-Zero).""" + block = AdaLNResBlock(dim=8, cond_dim=4) + x = torch.randn(5, 8) + cond = torch.randn(5, 4) + torch.testing.assert_close(block(x, cond), x) + + # ── EnergyRouter learn_width / learn_temperature ──────────────────────────── @@ -1077,6 +1128,44 @@ def test_build_models_explicit_resmlp_trunk_type_matches_default(): assert default_params == explicit_params +def test_build_models_explicit_add_block_conditioning_matches_default(): + """stage1_model.trunk.block_conditioning = 'add' is the default's + spelled-out equivalent, not a behaviour change — gitea #34.""" + default_cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2) + explicit_cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2, trunk={"type": "resmlp", "block_conditioning": "add"}) + default_stage1 = build_models(default_cfg)["stage1"] + explicit_stage1 = build_models(explicit_cfg)["stage1"] + assert default_stage1 is not None and explicit_stage1 is not None + assert type(default_stage1.trunk.blocks[0]) is type(explicit_stage1.trunk.blocks[0]) is ResBlock + default_params = sum(p.numel() for p in default_stage1.parameters()) + explicit_params = sum(p.numel() for p in explicit_stage1.parameters()) + assert default_params == explicit_params + + +@pytest.mark.parametrize("block_type,cls", [("film", FilmResBlock), ("adaln", AdaLNResBlock)]) +def test_build_models_selects_block_conditioning(block_type, cls): + cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2, trunk={"type": "resmlp", "block_conditioning": block_type}) + stage1 = build_models(cfg)["stage1"] + assert stage1 is not None + assert isinstance(stage1.trunk, ExpertTrunk) + assert all(isinstance(b, cls) for b in stage1.trunk.blocks) + + +def test_build_models_routed_trunk_uses_block_conditioning_for_every_expert(): + cfg = _nested_cfg( + pdg_vocab=4, + mat_vocab=2, + trunk={"type": "resmlp", "block_conditioning": "film"}, + stage1_router={"enabled": True, "type": "energy", "n_experts": 3}, + ) + stage1 = build_models(cfg)["stage1"] + assert stage1 is not None + assert isinstance(stage1.trunk, RoutedTrunk) + assert len(stage1.trunk.experts) == 3 + for expert in stage1.trunk.experts: + assert all(isinstance(b, FilmResBlock) for b in expert.blocks) + + def test_build_models_routed_pair_is_drop_in_for_sample_flow(): """Exercise the exact calling convention giant/sample.py uses.""" from giant.sample import sample_flow, sample_secondaries