noether.modeling.models.ab_upt

Attributes

KVPair

LayerCache

ModelKVCache

Classes

AnchoredBranchedUPT

Implementation of the Anchored Branched UPT model. Including input embedding and output projection, so this is an off-the-shelf model that can be used directly by providing the appropriate input tensors.

Module Contents

noether.modeling.models.ab_upt.KVPair
noether.modeling.models.ab_upt.LayerCache
noether.modeling.models.ab_upt.ModelKVCache
class noether.modeling.models.ab_upt.AnchoredBranchedUPT(config)

Bases: torch.nn.Module

Implementation of the Anchored Branched UPT model. Including input embedding and output projection, so this is an off-the-shelf model that can be used directly by providing the appropriate input tensors.

Parameters:

config (noether.core.schemas.models.AnchorBranchedUPTConfig) – Configuration for the AB-UPT model. See AnchorBranchedUPTConfig for details.

data_specs
rope
pos_embed
domain_names: list[str]
domain_biases
hidden_dim
physics_blocks
use_geometry_branch = False
domain_decoder_blocks
domain_decoder_projections
geometry_branch_forward(geometry_position, geometry_supernode_idx, geometry_batch_idx, condition, geometry_attn_kwargs)

Forward pass through the geometry branch of the model.

Parameters:
Return type:

torch.Tensor

physics_blocks_forward(domain_positions_all, geometry_encoding, physics_token_specs, physics_attn_kwargs, physics_perceiver_attn_kwargs, condition, kv_cache=None)

Forward pass through the physics blocks of the model.

Parameters:
Return type:

tuple[torch.Tensor, list[LayerCache]]

decoder_blocks_forward(x_physics, physics_token_specs, per_domain_token_specs, decoder_attn_kwargs, condition, kv_cache=None, domain_positions_all=None)

Forward pass through the per-domain decoder blocks.

Returns:

Tuple of (domain_predictions, new_domain_caches).

Parameters:
Return type:

tuple[dict[str, torch.Tensor], dict[str, list[LayerCache]]]

create_rope_frequencies(domain_positions_all, geometry_position=None, geometry_supernode_idx=None)

Create RoPE frequencies for all relevant positions.

Returns:

Tuple of (geometry_attn_kwargs, decoder_attn_kwargs, physics_perceiver_attn_kwargs, physics_attn_kwargs). decoder_attn_kwargs is keyed by domain name.

Parameters:
Return type:

tuple[dict[str, Any], dict[str, dict[str, Any]], dict[str, Any], dict[str, Any]]

forward(geometry_position=None, geometry_supernode_idx=None, geometry_batch_idx=None, domain_anchor_positions=None, domain_query_positions=None, domain_features=None, conditioning_inputs=None, kv_cache=None)

Forward pass of the AB-UPT model.

Example:

model(
    geometry_position=...,
    geometry_supernode_idx=...,
    geometry_batch_idx=...,
    domain_anchor_positions={"surface": surface_pos, "volume": volume_pos},
    domain_query_positions={"surface": query_pos},
    conditioning_inputs={"geometry_design_parameters": design_params},
)
Parameters:

  • geometry_position (torch.Tensor | None) – Coordinates of the geometry mesh. Tensor of shape (B * N_geometry, D_pos).

  • geometry_supernode_idx (torch.Tensor | None) – Supernode indices for the geometry points.

  • geometry_batch_idx (torch.Tensor | None) – Batch indices for the geometry points.

  • domain_anchor_positions (dict[str, torch.Tensor] | None) – Per-domain anchor positions, e.g. {"surface": (B, N, D), "volume": (B, M, D)}.

  • domain_query_positions (dict[str, torch.Tensor] | None) – Per-domain query positions (optional).

  • conditioning_inputs (dict[str, torch.Tensor] | None) – Conditioning tensors, e.g. {"geometry_design_parameters": (B, D)}.

  • kv_cache (ModelKVCache | None) – KV cache from a previous forward call.

  • domain_features (dict[str, torch.Tensor] | None)

Returns:

Tuple of (predictions, kv_cache).

Return type:

tuple[dict[str, torch.Tensor], ModelKVCache]