Core neural network and math
Reusable transformer shapes, patching, layers, normalization, gradients, schedulers, and camera geometry.
On this page
This group contains model-independent tensor operations and small modules. They encode shape contracts that otherwise drift between integrations: attention head width, MLP width, head split/merge, causal masks, image patch grids, normalization, stochastic depth, gradient clipping, rotation conversion, and diffusion scheduling.
Prefer reversible helpers for shape changes
patchify_image returns both token data and a PatchGridSpec. Keep the spec with the tokens; it records the original layout, batch dimensions, patch size, grid, and channel count needed for a checked inverse.
import torch
from worldfoundry.core import patchify_image, unpatchify_image
image = torch.arange(2 * 3 * 8 * 12).reshape(2, 3, 8, 12)
tokens, grid = patchify_image(image, patch_size=(4, 3), layout="nchw")
assert tokens.shape == (2, 8, 36)
restored = unpatchify_image(tokens, grid)
assert torch.equal(restored, image)unpatchify_image validates the exact token count and vector width instead of silently reshaping incompatible data. split_attention_heads and merge_attention_heads follow the same reversible principle for (..., sequence, hidden) tensors.
Derive dimensions in one place
from worldfoundry.core import transformer_shape_spec
shape = transformer_shape_spec(
hidden_size=1536,
num_heads=24,
mlp_ratio=8 / 3,
multiple_of=256,
)
assert shape.head_dim == 64
assert shape.mlp_hidden_size % 256 == 0Use these helpers in config validation and construction, not in a hot per-token loop. They reject non-positive or non-divisible dimensions early and keep model code from carrying slightly different rounding formulas.
Layers versus functional primitives
Mlp, SwiGLUFFN, PatchEmbed, DropPath, and LayerScale are reusable modules. Functional counterparts such as drop_path, rms_norm, layer_scale, and causal_attention_mask are useful when a model already owns parameters or needs a custom wrapper. Camera/rotation functions state their convention in the name; do not mix ZYX, OpenCV, and WXYZ quaternion conventions without an explicit conversion.
Complete reference
The blocks below are the generated signatures for this category. Use the on-page symbol index to jump; source links open the defining implementation behind each lazy export.
38 public symbols
attention_head_dim
def attention_head_dim(hidden_size: int,num_heads: int) -> intReturn per-head attention width and validate divisibility. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: int.
Parameters
num_headsintReturns
causal_attention_mask
def causal_attention_mask(query_len: int,key_len: int | None = None,include_self: bool = True) -> np.ndarrayOverview
Return a boolean causal mask for query/key lengths. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: np.ndarray.
Source docstring
Return a boolean causal mask for query/key lengths.
When `key_len is larger than query_len`, the mask assumes the query is aligned to the tail of a key/value cache.
Parameters
query_lenintkey_lenint | Nonedefault Noneinclude_selfbooldefault TrueReturns
clip_grad_norm_
def clip_grad_norm_(parameters: torch.Tensor | Iterable[torch.Tensor],max_norm: float,norm_type: float = 2.0,error_if_nonfinite: bool = False,foreach: bool | None = None,pp_mesh: torch.distributed.device_mesh.DeviceMesh | None = None) -> torch.TensorOverview
Clip gradients and optionally reduce the norm across pipeline-parallel stages. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: torch.Tensor.
Source docstring
Clip gradients and optionally reduce the norm across pipeline-parallel stages.
Same contract as `torch.nn.utils.clip_grad_norm_ plus *pp_mesh*: when set, the local total norm is all-reduced on that DeviceMesh (MAX for inf, SUM of norm**p then **(1/p)` otherwise) before scaling. DTensor norms are converted to a dense tensor first.
Parameters
parameterstorch.Tensor | Iterable[torch.Tensor]max_normfloatnorm_typefloatdefault 2.0error_if_nonfinitebooldefault Falseforeachbool | Nonedefault Nonepp_meshtorch.distributed.device_mesh.DeviceMesh | Nonedefault NoneReturns
clip_grads_with_norm_
def clip_grads_with_norm_(parameters: torch.Tensor | Iterable[torch.Tensor],max_norm: float,total_norm: torch.Tensor,foreach: bool | None = None) -> NoneOverview
Scale parameter gradients in-place using a precomputed total norm. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: None.
Source docstring
Scale parameter gradients in-place using a precomputed total norm.
The coefficient is `min(1, max_norm / (total_norm + 1e-6))`. Foreach multiply is used per device when available. No-op when no parameter has a gradient.
Parameters
parameterstorch.Tensor | Iterable[torch.Tensor]max_normfloattotal_normtorch.Tensorforeachbool | Nonedefault NoneReturns
def drop_path(x: Tensor,drop_prob: float = 0.0,training: bool = False,scale_by_keep: bool = True) -> TensorFunctional form of per-sample stochastic depth for residual branches.
Parameters
xTensordrop_probfloatdefault 0.0trainingbooldefault Falsescale_by_keepbooldefault TrueReturns
class DropPath(drop_prob: float | None = 0.0,scale_by_keep: bool = True)Stochastic depth module that drops residual paths per sample during training.
Parameters
drop_probfloat | Nonedefault 0.0scale_by_keepbooldefault TrueMethods
No-op in eval; in train, scale surviving samples when requested.
Parameters
xTensorReturns
def euler_angles_to_rotation_matrix_zyx(euler_angles: np.ndarray) -> np.ndarrayConvert `[x, y, z] Euler angles to Rz @ Ry @ Rx. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: np.ndarray`.
Parameters
euler_anglesnp.ndarrayReturns
class FlowMatchScheduler(num_inference_steps: int = 100,num_train_timesteps: int = 1000,shift: float = 3.0,sigma_max: float = 1.0,sigma_min: float = 0.003 / 1.002,inverse_timesteps: bool = False,extra_one_step: bool = False,reverse_sigmas: bool = False,exponential_shift: bool = False,exponential_shift_mu: float | None = None,shift_terminal: float | None = None)Overview
Flow-matching scheduler shared by bundled video runtimes. Belongs to Core neural-net/math helpers (shared tensor transforms).
Source docstring
Flow-matching scheduler shared by bundled video runtimes.
Why shift exists: a linear sigma grid spends too many steps at high noise for video. `shift * σ / (1 + (shift-1)σ) (or the exponential form) stretches the low-noise end. denoising_strength < 1 starts from an intermediate sigma for img2vid / inpaint. extra_one_step` drops the terminal sample so the last model call is still a real denoise step.
Parameters
num_inference_stepsintdefault 100num_train_timestepsintdefault 1000shiftfloatdefault 3.0sigma_maxfloatdefault 1.0sigma_minfloatdefault 0.003 / 1.002inverse_timestepsbooldefault Falseextra_one_stepbooldefault Falsereverse_sigmasbooldefault Falseexponential_shiftbooldefault Falseexponential_shift_mufloat | Nonedefault Noneshift_terminalfloat | Nonedefault NoneMethods
set_timesteps(num_inference_steps: int = 100,denoising_strength: float = 1.0,training: bool = False,shift: float | None = None,dynamic_shift_len: int | None = None,exponential_shift_mu: float | None = None) -> NoneSourceRebuild `sigmas / timesteps` for a new step count or shift.
Parameters
num_inference_stepsintdefault 100denoising_strengthfloatdefault 1.0trainingbooldefault Falseshiftfloat | Nonedefault Nonedynamic_shift_lenint | Nonedefault Noneexponential_shift_mufloat | Nonedefault NoneReturns
Overview
One Euler step: `sample + velocity * (σ_next - σ)`.
Source docstring
One Euler step: `sample + velocity * (σ_next - σ)`.
`to_final forces the terminal sigma even mid-schedule (last-step snap). Inverse/reversed calendars use terminal 1.0 instead of 0.0` so the integrator still walks toward clean data.
Parameters
model_outputtimestepsampleto_finalbooldefault FalseRecover a velocity that maps `sample to stabilized_sample`.
Parameters
timestepsamplestabilized_sampleRun the forward corruption process.
Parameters
original_samplesnoisetimestepFlow-matching target `noise - sample` (velocity toward noise).
Parameters
samplenoisetimestepPer-timestep loss weight from the Gaussian bump built in training mode.
Parameters
timestepcalculate_shift(image_seq_len: int,base_seq_len: int = 256,max_seq_len: int = 8192,base_shift: float = 0.5,max_shift: float = 0.9) -> floatSourceInterpolate exponential-shift μ from sequence length (Flux-style).
Parameters
image_seq_lenintbase_seq_lenintdefault 256max_seq_lenintdefault 8192base_shiftfloatdefault 0.5max_shiftfloatdefault 0.9Returns
get_total_norm
def get_total_norm(tensors: torch.Tensor | Iterable[torch.Tensor],norm_type: float = 2.0,error_if_nonfinite: bool = False,foreach: bool | None = None) -> torch.TensorOverview
Compute the total norm of tensors as if their flattened values were concatenated. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: torch.Tensor.
Source docstring
Compute the total norm of tensors as if their flattened values were concatenated.
Groups tensors by device/dtype and uses `torch._foreach_norm when the device supports it (or *foreach* is forced). Empty input yields a CPU scalar 0.0`. Set *error_if_nonfinite* to raise instead of returning NaN/Inf.
Parameters
tensorstorch.Tensor | Iterable[torch.Tensor]norm_typefloatdefault 2.0error_if_nonfinitebooldefault Falseforeachbool | Nonedefault NoneReturns
layer_scale
def layer_scale(value: Any,scale: Any,bias: Any = None) -> AnyApply a stateless elementwise layer-scale transform. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: Any.
Parameters
valueAnyscaleAnybiasAnydefault NoneReturns
LayerScale
class LayerScale(dim: int,init_values: Union[float, Tensor] = 1e-05,inplace: bool = False,device = None)Learnable per-channel residual scaling used in modern Vision Transformer blocks.
Parameters
dimintinit_valuesUnion[float, Tensor]default 1e-05gamma`.inplacebooldefault Falseforward`.devicedefault NoneMethods
Copy a tensor init or fill a scalar; used after meta-device materialize.
Returns
Scale the final dimension of `x by the learned gamma`.
Parameters
xTensorReturns
Include init and inplace flags so module printouts distinguish LayerScales.
Returns
merge_attention_heads
def merge_attention_heads(value: Any) -> AnyInvert `split_attention_heads for (..., heads, seq, head_dim) tensors. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: Any`.
Parameters
valueAnyReturns
class Mlp(in_features: int,hidden_features: Optional[int] = None,out_features: Optional[int] = None,act_layer: Callable[..., nn.Module] = nn.GELU,drop: float | tuple[float, float] = 0.0,bias: bool | tuple[bool, bool] = True,device = None)timm-style 2-layer ViT FFN (fc1 → act → fc2); distinct from `SamHeadMLP`. Belongs to Core neural-net/math helpers (shared tensor transforms).
Parameters
in_featuresintout_featuresOptional[int]default Nonein_features`.act_layerCallable[..., nn.Module]default nn.GELUdropfloat | tuple[float, float]default 0.0biasbool | tuple[bool, bool]default Truedevicedefault NoneMethods
Apply `fc1 → activation → dropout → fc2 → dropout`.
Parameters
xTensorReturns
named_apply
def named_apply(fn: Callable[..., object],module: nn.Module,name: str = '',depth_first: bool = True,include_root: bool = False) -> nn.ModuleOverview
Apply `fn(module=..., name=...) recursively with stable module names. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: nn.Module`.
Source docstring
Apply `fn(module=..., name=...)` recursively with stable module names.
This is the named equivalent of :meth:torch.nn.Module.apply. By default, children are visited depth-first and the supplied root is omitted, matching the behavior historically used by the DINO-style backbones in WorldFoundry.
Parameters
fnCallable[..., object]modulenn.Modulenamestrdefault ''depth_firstbooldefault Trueinclude_rootbooldefault FalseReturns
def named_camera_trajectory_tensors(trajectory: str,initial_world_to_camera: Any,initial_intrinsic: Any,num_frames: int,movement_distance: float = 0.3,camera_rotation: str = 'center_facing',center_depth: float = 1.0,num_circles: int = 1)Overview
Build the eight named GEN3C camera trajectories without a model runtime. Belongs to Core neural-net/math helpers (shared tensor transforms).
Source docstring
Build the eight named GEN3C camera trajectories without a model runtime.
The output is `(world_to_camera, intrinsics) with batched shapes [1,T,4,4] and [1,T,3,3]. Keeping this geometry in core` makes it reusable by camera-conditioned recipes instead of tying it to Cosmos.
Parameters
trajectorystrinitial_world_to_cameraAnyinitial_intrinsicAnynum_framesintmovement_distancefloatdefault 0.3camera_rotationstrdefault 'center_facing'center_depthfloatdefault 1.0num_circlesintdefault 1PatchEmbed
class PatchEmbed(img_size: Union[int, tuple[int, int]] = 224,patch_size: Union[int, tuple[int, int]] = 16,in_chans: int = 3,embed_dim: int = 768,norm_layer: Optional[Callable[..., nn.Module]] = None,flatten_embedding: bool = True)2D image patch embedding: `(B, C, H, W) -> (B, N, D)`. Belongs to Core neural-net/math helpers (shared tensor transforms).
Parameters
img_sizeUnion[int, tuple[int, int]]default 224patch_sizeUnion[int, tuple[int, int]]default 16in_chansintdefault 3embed_dimintdefault 768norm_layerOptional[Callable[..., nn.Module]]default Noneflatten_embeddingbooldefault True(B, N, D) when true, otherwise (B, grid_h, grid_w, D)`.Methods
Uniform init scaled by fan-in of the patch convolution (checkpoint match).
Returns
Project a divisible NCHW image batch into patch embeddings.
Parameters
xTensorReturns
Raises
ValueError- Runtime height or width is not divisible by patch size.
Report conv + optional norm FLOPs for the *nominal* `img_size` grid.
Returns
PatchEmbed_Mlp
class PatchEmbed_Mlp(img_size: Union[int, tuple[int, int]] = 224,patch_size: Union[int, tuple[int, int]] = 16,in_chans: int = 3,embed_dim: int = 768,norm_layer: Optional[Callable[..., nn.Module]] = None,flatten_embedding: bool = True)Patch embedding implemented with pixel unshuffle and MLP projection. Belongs to Core neural-net/math helpers (shared tensor transforms).
Parameters
img_sizeUnion[int, tuple[int, int]]default 224patch_sizeUnion[int, tuple[int, int]]default 16in_chansintdefault 3embed_dimintdefault 768norm_layerOptional[Callable[..., nn.Module]]default Noneflatten_embeddingbooldefault TruePatchGridSpec
class PatchGridSpec(original_shape: tuple[int, ...],patch_size: tuple[int, int],layout: ImageLayout = 'nchw')Shape contract needed to invert a 2D image patchification. Belongs to Core neural-net/math helpers (shared tensor transforms).
Attributes
original_shapetuple[int, ...]patch_sizetuple[int, int]layoutImageLayoutdefault 'nchw'Methods
Leading axes before the CHW / HWC triple (may be empty).
Parameters
selfReturns
Channel axis selected by `layout` (not always dim 1).
Parameters
selfReturns
`(H, W)` regardless of NCHW vs NHWC.
Parameters
selfReturns
Number of patches along height and width.
Parameters
selfReturns
`grid_h * grid_w` — the sequence length after patchify.
Parameters
selfReturns
patchify_image
def patchify_image(value: Any,patch_size: int | Sequence[int],layout: ImageLayout = 'nchw') -> tuple[Any, PatchGridSpec]Overview
Convert an image tensor into flattened 2D patch tokens. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: tuple[Any, PatchGridSpec].
Source docstring
Convert an image tensor into flattened 2D patch tokens.
`layout="nchw" expects (..., C, H, W) and layout="nhwc" expects (..., H, W, C). The returned token tensor has shape (..., grid_h * grid_w, C * patch_h * patch_w)`.
Parameters
valueAnypatch_sizeint | Sequence[int]layoutImageLayoutdefault 'nchw'Returns
class Permute(dims: tuple[int, ])Module wrapper around `Tensor.permute`. Belongs to Core neural-net/math helpers (shared tensor transforms).
Attributes
dimstuple[int, ...]Methods
Apply the stored axis permutation; no shape validation (caller contract).
Parameters
valueTensorReturns
PixelUnshuffle
class PixelUnshuffle(downscale_factor: int)Module wrapper for `torch.nn.functional.pixel_unshuffle`. Belongs to Core neural-net/math helpers (shared tensor transforms).
Parameters
downscale_factorintMethods
Unshuffle `[B, C, H, W]; empty tensors use view` because F.
Parameters
valueTensorReturns
Raises
ValueError- empty input spatial dims are 0 or not divisible by the factor.
def quaternion_xyzw_to_rotation_matrix(quaternions: Tensor) -> TensorConvert `[..., x, y, z, w] quaternions to rotation matrices. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: Tensor`.
Parameters
quaternionsTensorReturns
ray_condition
def ray_condition(K,c2w,H: int,W: int,device,flip_flag = None,use_ray_o: bool = False)Build per-pixel Plucker ray features from intrinsics and camera-to-world poses. Belongs to Core neural-net/math helpers (shared tensor transforms).
Parameters
K[fx, fy, cx, cy] layout with shape [B,V,4]. When None`, rays use a constant forward camera-space direction.c2w[B,V,4,4]`.HintWintdeviceflip_flagdefault Noneuse_ray_obooldefault False[ray_origin, ray_direction] instead of the Plucker [ray_origin x ray_direction, ray_direction]` form.def rms_norm(value: Any,weight: Any = None,eps: float = 1e-06) -> AnyApply RMS normalization over the last dimension. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: Any.
Parameters
valueAnyweightAnydefault Noneepsfloatdefault 1e-06Returns
def rotation_matrix_to_euler_angles_opencv(rotation: np.ndarray) -> tuple[float, float, float]Return the ZYX angles in OpenCV's `(z, y, x) order. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: tuple[float, float, float]`.
Parameters
rotationnp.ndarrayReturns
def rotation_matrix_to_euler_angles_zyx(rotation: np.ndarray) -> np.ndarrayConvert a 3x3 rotation matrix to `[x, y, z] Euler angles. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: np.ndarray`.
Parameters
rotationnp.ndarrayReturns
def rotation_matrix_to_quaternion_wxyz(rotation: np.ndarray) -> np.ndarrayConvert a 3x3 rotation matrix to a normalized `[w, x, y, z] quaternion. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: np.ndarray`.
Parameters
rotationnp.ndarrayReturns
def rotation_matrix_to_quaternion_xyzw(matrix: Tensor) -> TensorOverview
Convert `[..., 3, 3] rotation matrices to canonical XYZW quaternions. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: Tensor`.
Source docstring
Convert `[..., 3, 3]` rotation matrices to canonical XYZW quaternions.
The implementation chooses the best-conditioned of four equivalent quaternion candidates and keeps gradients well-defined away from zero.
Parameters
matrixTensorReturns
class SchedulerInterface(alphas_cumprod: torch.Tensor)Base interface for diffusion noise schedules. Belongs to Core neural-net/math helpers (shared tensor transforms).
Attributes
alphas_cumprodtorch.TensorMethods
Run the forward corruption process.
Parameters
clean_latenttorch.Tensornoisetorch.Tensortimesteptorch.Tensorconvert_x0_to_noise(x0: torch.Tensor,xt: torch.Tensor,timestep: torch.Tensor) -> torch.TensorSourceConvert a clean-data prediction to a noise prediction.
Parameters
x0torch.Tensorxttorch.Tensortimesteptorch.TensorReturns
convert_noise_to_x0(noise: torch.Tensor,xt: torch.Tensor,timestep: torch.Tensor) -> torch.TensorSourceConvert a noise prediction to a clean-data prediction.
Parameters
noisetorch.Tensorxttorch.Tensortimesteptorch.TensorReturns
convert_velocity_to_x0(velocity: torch.Tensor,xt: torch.Tensor,timestep: torch.Tensor) -> torch.TensorSourceConvert a velocity prediction to a clean-data prediction.
Parameters
velocitytorch.Tensorxttorch.Tensortimesteptorch.TensorReturns
split_attention_heads
def split_attention_heads(value: Any,num_heads: int) -> AnyReshape `(..., seq, hidden) to (..., heads, seq, head_dim). Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: Any`.
Parameters
valueAnynum_headsintReturns
def standardize_quaternion_xyzw(quaternions: Tensor) -> TensorChoose the equivalent XYZW quaternion whose real part is non-negative. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: Tensor.
Parameters
quaternionsTensorReturns
class SwiGLUFFN(in_features: int,hidden_features: Optional[int] = None,out_features: Optional[int] = None,act_layer: Callable[..., nn.Module] | None = None,drop: float = 0.0,bias: bool = True)SwiGLU feed-forward layer shared across model integrations.
Parameters
in_featuresintout_featuresOptional[int]default Noneact_layerCallable[..., nn.Module] | Nonedefault Nonedropfloatdefault 0.0biasbooldefault TrueMethods
Last-dim SwiGLU matching the `w12 / w3` checkpoint split.
Parameters
xTensorReturns
SwiGLUFFNFused
class SwiGLUFFNFused(in_features: int,hidden_features: Optional[int] = None,out_features: Optional[int] = None,act_layer: Callable[..., nn.Module] | None = None,drop: float = 0.0,bias: bool = True)SwiGLU FFN with hidden width rounded for tensor-core-friendly matmuls. Belongs to Core neural-net/math helpers (shared tensor transforms).
Parameters
in_featuresintout_featuresOptional[int]default Nonein_features`.act_layerCallable[..., nn.Module] | Nonedefault Nonedropfloatdefault 0.0biasbooldefault Truetransformer_shape_spec
def transformer_shape_spec(hidden_size: int,num_heads: int,mlp_ratio: float = 4.0,multiple_of: int | None = None) -> TransformerShapeSpecBuild a reusable shape spec for attention and MLP dimensions. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: TransformerShapeSpec.
Parameters
num_headsintmlp_ratiofloatdefault 4.0multiple_ofint | Nonedefault NoneReturns
class TransformerShapeSpec(hidden_size: int,num_heads: int,head_dim: int,mlp_hidden_size: int)Common transformer dimensions derived from hidden size and head count. Belongs to Core neural-net/math helpers (shared tensor transforms).
Attributes
num_headsinthead_dimintunpatchify_image
def unpatchify_image(patches: Any,spec: PatchGridSpec) -> AnyInvert `patchify_image using the returned PatchGridSpec. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: Any`.
Parameters
patchesAnyspecPatchGridSpecReturns
zero_module
def zero_module(module: nn.Module) -> nn.ModuleDetach and zero all parameters in a module. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: nn.Module.
Parameters
modulenn.Module