gensbi.recipes.utils

Attributes

_EMBEDDINGS_1D
_EMBEDDINGS_2D
_MOVED_TO_PATCHING

Functions

__getattr__(name)
_normalize_patch_size(size)	Normalize a patch-size spec into an (obs, cond) tuple.
_require_channel(x[, name])	Enforce a tabular channel axis (B, dim, C); reject a bare (B, dim).
_resolve_embedding_ids(dim, strategy, semantic_id[, size])	Resolve ID embeddings by strategy name.
_single_obs(x_o, *, channel[, name])	Canonicalize a single conditioning observation, then enforce batch == 1.
_validate_base_pixels(base_pixels)	Validate a base-pixel subset spec; return it as a list of ints.
build_edm_path(sde, config)	Build an EDM-family diffusion path from an SDE type string and config.
build_sm_path(sde_type, config)	Build a score-matching path from an SDE type string and config.
healpix_rope_theta(nside)	Suggested RoPE theta for a full-sky HEALPix token grid at nside.
init_ids_1d(dim[, semantic_id])	Build 1D positional IDs, returning (ids, dim).
init_ids_2d(dim[, semantic_id, size])	Build 2D positional IDs for a patchified image grid.
init_ids_healpix(nside[, base_pixels])	Build spherical RoPE ids for tokens on a HEALPix grid, returning
init_ids_joint(dim_obs, dim_cond)
parse_training_config(config_path)	Parse training and optimizer configuration from a YAML config file.
scale_lr(batch_size[, base_lr, reference_batch_size])	Scale learning rate based on batch size using square root scaling.

Module Contents

gensbi.recipes.utils.__getattr__(name)

gensbi.recipes.utils._normalize_patch_size(size)

Normalize a patch-size spec into an (obs, cond) tuple.

Parameters:

size (int or tuple of int) – A single int is broadcast to both inputs (8 -> (8, 8)). A length-2 tuple is taken as (obs_size, cond_size) so the two inputs can use different patch sizes. Use 1 for an input that is not patchified.

Returns:

(obs_size, cond_size).

Return type:

tuple of int

gensbi.recipes.utils._require_channel(x, name='input')

Enforce a tabular channel axis (B, dim, C); reject a bare (B, dim).

gensbi.recipes.utils._resolve_embedding_ids(dim, strategy, semantic_id, size=2)

Resolve ID embeddings by strategy name.

Parameters:

• dim (int or tuple of int) – Dimension specification (number of tokens, or (H, W) for 2D images).

• strategy (str or IdStrategy) – Embedding strategy name (e.g., “absolute”, “pos1d”, “rope1d”, “pos2d”, “rope2d”) or a strategy object with build(dim) (e.g. gensbi.recipes.HealpixRope). NOTE these pipeline-side builder names are a different vocabulary from the model-side id_embedding_strategy strings (where “rope” means “apply RoPE to the ids the pipeline provides”) — see gensbi.recipes.id_strategies.

• semantic_id (int) – Semantic identifier for the token group (0=obs, 1=cond).

• size (int, optional) – Patch edge length for 2D strategies (default 2). Ignored for 1D strategies. Use 1 for no patchification.

Returns:

• ids (Array) – Token ID array.

• resolved_dim (int) – Resolved flat dimension.

Raises:

ValueError – If strategy is not recognized.

gensbi.recipes.utils._single_obs(x_o, *, channel, name='x_o')

Canonicalize a single conditioning observation, then enforce batch == 1.

Shape handling comes FIRST so a misshaped input can never be misread as a batch (e.g. (dim, C) read as dim observations):

• channel="require": tabular flow-pipeline contract — input must already carry batch and channel axes (1, dim, C); channel-less input raises ValueError (same _require_channel() as training).

• channel="promote": FM-pipeline contract — lenient promotion: (dim,) -> (1, dim, 1) and (B, dim) -> (B, dim, 1).

• channel="none": structured inputs — the model owns the trailing shape; only a leading batch axis (ndim >= 2) is required.

A leading batch axis > 1 then raises ValueError: single-observation methods never silently discard observations — use sample_batched. Returns the canonicalized array with its size-1 batch axis kept.

gensbi.recipes.utils._validate_base_pixels(base_pixels)

Validate a base-pixel subset spec; return it as a list of ints.

gensbi.recipes.utils.build_edm_path(sde, config)

Build an EDM-family diffusion path from an SDE type string and config.

Parameters:

• sde (str) – SDE type: "EDM", "VE", or "VP".

• config (dict) – Training configuration dict; scheduler hyperparameters are read from here with sensible defaults.

Returns:

Configured diffusion path.

Return type:

EDMPath

Raises:

ValueError – If sde is not one of {"EDM", "VE", "VP"}.

gensbi.recipes.utils.build_sm_path(sde_type, config)

Build a score-matching path from an SDE type string and config.

Parameters:

• sde_type (str) – SDE type: "VP" or "VE".

• config (dict) – Training configuration dict; scheduler hyperparameters are read from here with sensible defaults.

Returns:

Configured score-matching path.

Return type:

SMPath

Raises:

ValueError – If sde_type is not one of {"VP", "VE"}.

gensbi.recipes.utils.healpix_rope_theta(nside)

Suggested RoPE theta for a full-sky HEALPix token grid at nside.

Follows the project convention theta = 10 * token count (the same rule Flux1Params applies by default via 10 * (dim_obs + dim_cond)): a full-sky grid has 12 * nside**2 tokens. Exposed so spherical models can derive theta from the intuitive knob (nside, always known after the encoder) instead of setting it by hand.

Parameters:

nside (int)

Return type:

int

gensbi.recipes.utils.init_ids_1d(dim, semantic_id=None)

Build 1D positional IDs, returning (ids, dim).

ids is (1, dim, 1) when semantic_id is None (position only), or (1, dim, 2) otherwise, with the position at axis 0 and the semantic id at axis 1.

FIXME (axis-order footgun): this places the semantic id on the LAST axis, which is the reverse of init_ids_2d() (semantic at axis 0, then h, w – the established convention). The two are safe in isolation, but they do NOT line up if 1D and 2D ids are ever fed into the same RoPE grid (e.g. FieldDiT Phase-2 obs+cond co-tokenization with a shared EmbedND, where axes_dim[i] is matched to ids axis i). This should be unified to the 2D convention (semantic at axis 0). It is not changed here because callers that pass semantic_id – notably the Flux1 rope1d path via init_ids() – and any code indexing ids[..., k] must be updated in lockstep.

Parameters:

• dim (int)

• semantic_id (Union[int, None])

gensbi.recipes.utils.init_ids_2d(dim, semantic_id=0, size=2)

Build 2D positional IDs for a patchified image grid.

The grid has one entry per patch, i.e. (dim[0] // size, dim[1] // size), matching patchify_2d(x, size=size). size is the patch edge length; use size=1 for no patchification (one token per pixel).

Parameters:

• dim (Tuple[int, int])

• semantic_id (int)

• size (int)

gensbi.recipes.utils.init_ids_healpix(nside, base_pixels=None)

Build spherical RoPE ids for tokens on a HEALPix grid, returning (ids, num_tokens).

Method: standard N-dimensional RoPE (RoFormer, arXiv:2104.09864 — the mechanism implemented by Flux1’s EmbedND) applied uniformly, on all three axes and all frequency bands, to the 3D Cartesian coordinates of HEALPix pixel centers on the unit sphere. Each token maps to its pixel-center unit vector (healpy.pix2vec, NEST ordering), scaled to pixel units so adjacent tokens differ by ~1 in coordinate (radius nside * sqrt(3/pi) = 1/pixel angular size), which keeps theta’s semantics identical to 2D-image usage (see healpix_rope_theta()).

Attention scores then depend on positions only through the chord vector n_q - n_k, whose norm 2 sin(gamma/2) is strictly monotone in great-circle distance gamma — geodesic geometry with no projection step, hence no face-seam or polar artifacts, and any base_pixels subset works by construction. Caveat: d(chord)/d(gamma) -> 0 at antipodes, so resolution among near-antipodal separations is mildly compressed (benign for near/far attention). This is NOT an adaptation of SpheRoPE (arXiv:2606.32033 — closest prior work; ERP grid, pretrained constraints); see also StereoRoPE (arXiv:2606.31248, documents the failure of index-based RoPE on HEALPix) and Unlu (arXiv:2310.04454, an SO(3) feature-rotation alternative not adopted). Full rationale: docs/superpowers/specs/2026-07-19-healpix-rope-design.md.

Use with Flux1 via id_embedding_strategy=("absolute", "rope") and a 3-entry axes_dim (each even, summing to the per-head dim, e.g. (22, 22, 20) for 64). Obs/theta-stream tokens automatically get origin (0, 0, 0) rope ids — the identity rotation, i.e. an exactly isotropic positional readout of the conditioning tokens.

Parameters:

• nside (int) – HEALPix resolution of the token grid (power of 2). With HEAL-SWIN style encoders this is the bottleneck nside; tokens must correspond to single HEALPix pixels (power-of-4 pixels-per-token upstream).

• base_pixels (sequence of int, optional) – Base pixels (0..11) covered by the token grid, for partial-sky models. None (default) means full sky. Tokens are ordered by base pixel as given, NEST within each.

Returns:

• ids (jax.Array) – (1, num_tokens, 3) float32 scaled pixel-center coordinates.

• num_tokens (int) – len(base_pixels) * nside**2.

gensbi.recipes.utils.init_ids_joint(dim_obs, dim_cond)

Parameters:

• dim_obs (int)

• dim_cond (int)

gensbi.recipes.utils.parse_training_config(config_path)

Parse training and optimizer configuration from a YAML config file.

Reads the training and optimizer sections of the config and returns a flat dictionary consumed by AbstractPipeline.

Parameters:

config_path (str) – Path to the YAML configuration file.

Returns:

training_config – Parsed training configuration dictionary.

Return type:

dict

gensbi.recipes.utils.scale_lr(batch_size, base_lr=0.0001, reference_batch_size=256)

Scale learning rate based on batch size using square root scaling.

Parameters:

• batch_size (int) – The current batch size.

• base_lr (float) – The base learning rate for the reference batch size.

• reference_batch_size (int, optional) – The reference batch size. Defaults to 256.

Returns:

The adjusted learning rate.

Return type:

float

gensbi.recipes.utils._EMBEDDINGS_1D

gensbi.recipes.utils._EMBEDDINGS_2D

gensbi.recipes.utils._MOVED_TO_PATCHING = ('patchify_2d', 'depatchify_2d')