gensbi.recipes.utils#

Attributes#

_EMBEDDINGS_1D

_EMBEDDINGS_2D

Functions#

_resolve_embedding_ids(dim, strategy, semantic_id)

Resolve ID embeddings by strategy name.

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.

init_ids_1d(dim[, semantic_id])

init_ids_2d(dim[, semantic_id])

init_ids_joint(dim_obs, dim_cond)

parse_training_config(config_path)

Parse training and optimizer configuration from a YAML config file.

patchify_2d(x)

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._resolve_embedding_ids(dim, strategy, semantic_id)[source]#

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) – Embedding strategy name (e.g., “absolute”, “pos1d”, “rope1d”, “pos2d”, “rope2d”).

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

Returns:

  • ids (Array) – Token ID array.

  • resolved_dim (int) – Resolved flat dimension.

Raises:

ValueError – If strategy is not recognized.

gensbi.recipes.utils.build_edm_path(sde, config)[source]#

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)[source]#

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.init_ids_1d(dim, semantic_id=None)[source]#
Parameters:

  • dim (int)

  • semantic_id (Union[int, None])

gensbi.recipes.utils.init_ids_2d(dim, semantic_id=0)[source]#
Parameters:

  • dim (Tuple[int, int])

  • semantic_id (int)

gensbi.recipes.utils.init_ids_joint(dim_obs, dim_cond)[source]#
Parameters:

  • dim_obs (int)

  • dim_cond (int)

gensbi.recipes.utils.parse_training_config(config_path)[source]#

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.patchify_2d(x)[source]#
Parameters:

x (jax.Array)

gensbi.recipes.utils.scale_lr(batch_size, base_lr=0.0001, reference_batch_size=256)[source]#

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[source]#
gensbi.recipes.utils._EMBEDDINGS_2D[source]#