Core 神经网络与数学
可复用 Transformer 形状、patch、layer、normalization、gradient、scheduler 与相机几何。
这一组收录与模型身份无关的张量操作和小型 module。它们固化那些很容易在不同接入中漂移的形状契约：attention head 宽度、MLP 宽度、head 拆分合并、causal mask、图像 patch grid、normalization、stochastic depth、gradient clipping、旋转转换和 diffusion scheduler。
形状变换优先使用可逆 helper
patchify_image 同时返回 token 数据和 PatchGridSpec。应让 spec 与 token 一起流转，因为它记录了经过检查的逆变换所需的原始布局、batch 维、patch size、grid 和 channel 数。
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 会验证准确的 token 数和向量宽度，不会把不兼容数据静默 reshape。split_attention_heads 与 merge_attention_heads 对 (..., sequence, hidden) tensor 采用同样的可逆原则。
在一个地方推导维度
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 == 0这些 helper 适合用于配置验证与构造，不要放进每 token 的热循环。它们会提前拒绝非正数或不可整除的维度，并避免不同模型代码各自携带略有差异的 rounding 公式。
Module 与函数式 primitive
Mlp、SwiGLUFFN、PatchEmbed、DropPath 和 LayerScale 是可复用 module。当模型已经拥有参数或者需要自定义 wrapper 时，可以使用 drop_path、rms_norm、layer_scale 和 causal_attention_mask 等函数式入口。相机与旋转函数会在名称中声明 convention；不要在没有显式转换的情况下混用 ZYX、OpenCV 和 WXYZ quaternion convention。
完整参考
以下为该类别的生成签名。可用本页符号索引跳转；源码链接指向各惰性导出背后的具体实现。
37 个公开符号
def attention_head_dim(hidden_size: int, num_heads: int) -> intworldfoundry.core.attention_head_dimfrom worldfoundry.core import attention_head_dim简介
attention_head_dim — Return per-head attention width and validate divisibility. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：int。
参数
num_headsint
返回值: int
def causal_attention_mask(query_len: int,key_len: int | None = None,include_self: bool = True) -> np.ndarrayworldfoundry.core.causal_attention_maskfrom worldfoundry.core import causal_attention_mask简介
causal_attention_mask — Return a boolean causal mask for query/key lengths. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：np.ndarray。
源码 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.
参数
query_lenintkey_lenint | None- 默认值:
None include_selfbool- 默认值:
True
返回值: np.ndarray
clip_grad_norm_
funcdef 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.Tensorworldfoundry.core.clip_grad_norm_from worldfoundry.core import clip_grad_norm_简介
clip_grad_norm_ — Clip gradients and optionally reduce the norm across pipeline-parallel stages. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：torch.Tensor。
参数
parameterstorch.Tensor | Iterable[torch.Tensor]max_normfloatnorm_typefloat- 默认值:
2.0 error_if_nonfinitebool- 默认值:
False foreachbool | None- 默认值:
None pp_meshtorch.distributed.device_mesh.DeviceMesh | None- 默认值:
None
返回值: torch.Tensor
def clip_grads_with_norm_(parameters: torch.Tensor | Iterable[torch.Tensor],max_norm: float,total_norm: torch.Tensor,foreach: bool | None = None) -> Noneworldfoundry.core.clip_grads_with_norm_from worldfoundry.core import clip_grads_with_norm_简介
clip_grads_with_norm_ — Scale parameter gradients in-place using a precomputed total norm. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：None。
参数
parameterstorch.Tensor | Iterable[torch.Tensor]max_normfloattotal_normtorch.Tensorforeachbool | None- 默认值:
None
返回值: None
drop_path
funcdef drop_path(x: Tensor,drop_prob: float = 0.0,training: bool = False,scale_by_keep: bool = True) -> Tensorworldfoundry.core.drop_pathfrom worldfoundry.core import drop_path简介
残差分支上按样本 Stochastic Depth 的函数式接口。
参数
xTensordrop_probfloat- 默认值:
0.0 trainingbool- 默认值:
False scale_by_keepbool- 默认值:
True
返回值: Tensor
def euler_angles_to_rotation_matrix_zyx(euler_angles: np.ndarray) -> np.ndarrayworldfoundry.core.euler_angles_to_rotation_matrix_zyxfrom worldfoundry.core import euler_angles_to_rotation_matrix_zyx简介
euler_angles_to_rotation_matrix_zyx — Convert `[x, y, z] Euler angles to Rz @ Ry @ Rx. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：np.ndarray`。
参数
euler_anglesnp.ndarray
返回值: np.ndarray
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)worldfoundry.core.FlowMatchSchedulerfrom worldfoundry.core import FlowMatchScheduler简介
FlowMatchScheduler — Flow-matching scheduler shared by bundled video runtimes. 属于 Core 神经网络/数学辅助（共享张量变换）。
参数
num_inference_stepsint- 默认值:
100 num_train_timestepsint- 默认值:
1000 shiftfloat- 默认值:
3.0 sigma_maxfloat- 默认值:
1.0 sigma_minfloat- 默认值:
0.003 / 1.002 inverse_timestepsbool- 默认值:
False extra_one_stepbool- 默认值:
False reverse_sigmasbool- 默认值:
False
方法
set_timesteps(num_inference_steps: int = 100,denoising_strength: float = 1.0,training: bool = False) -> None源码简介
该类型上的公开 method。
参数
num_inference_stepsint- 默认值:
100 denoising_strengthfloat- 默认值:
1.0 trainingbool- 默认值:
False
返回值: None
简介
该类型上的公开 method。
参数
model_outputtimestepsampleto_finalbool- 默认值:
False
简介
add_noise — Run the forward corruption process.
参数
original_samplesnoisetimestep
get_total_norm
funcdef get_total_norm(tensors: torch.Tensor | Iterable[torch.Tensor],norm_type: float = 2.0,error_if_nonfinite: bool = False,foreach: bool | None = None) -> torch.Tensorworldfoundry.core.get_total_normfrom worldfoundry.core import get_total_norm简介
get_total_norm — Compute the total norm of tensors as if their flattened values were concatenated. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：torch.Tensor。
参数
tensorstorch.Tensor | Iterable[torch.Tensor]norm_typefloat- 默认值:
2.0 error_if_nonfinitebool- 默认值:
False foreachbool | None- 默认值:
None
返回值: torch.Tensor
layer_scale
funcdef layer_scale(value: Any,scale: Any,bias: Any = None) -> Anyworldfoundry.core.layer_scalefrom worldfoundry.core import layer_scale简介
layer_scale — Apply a stateless elementwise layer-scale transform. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：Any。
参数
valueAnyscaleAnybiasAny- 默认值:
None
返回值: Any
LayerScale
clsclass LayerScale(dim: int,init_values: Union[float, Tensor] = 1e-05,inplace: bool = False,device = None)worldfoundry.core.LayerScalefrom worldfoundry.core import LayerScale简介
现代 ViT block 中常见的可学习逐通道残差缩放。
参数
dimint- Channel width of the final tensor dimension.
init_valuesUnion[float, Tensor]- Scalar or per-channel initialization for `
gamma`.默认值:1e-05 inplacebool- Multiply the input in place during `
forward`.默认值:False device- Optional parameter device.默认值:
None
def merge_attention_heads(value: Any) -> Anyworldfoundry.core.merge_attention_headsfrom worldfoundry.core import merge_attention_heads简介
merge_attention_heads — Invert `split_attention_heads for (..., heads, seq, head_dim) tensors. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：Any`。
参数
valueAny
返回值: Any
Mlp
clsclass 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)worldfoundry.core.Mlpfrom worldfoundry.core import Mlp简介
Mlp — timm-style 2-layer ViT FFN (fc1 → act → fc2); distinct from `SamHeadMLP`. 属于 Core 神经网络/数学辅助（共享张量变换）。
参数
in_featuresint- Input width.
out_featuresOptional[int]- Output width; defaults to `
in_features`.默认值:None act_layerCallable[..., nn.Module]- Activation module factory between projections.默认值:
nn.GELU dropfloat | tuple[float, float]- One probability for both dropout sites or a pair for the first and second sites.默认值:
0.0 biasbool | tuple[bool, bool]- One bias flag for both projections or a pair.默认值:
True device- Optional projection parameter device.默认值:
None
方法
简介
forward — Apply `fc1 → activation → dropout → fc2 → dropout`.
参数
xTensor
返回值: Tensor
named_apply
funcdef named_apply(fn: Callable[..., object],module: nn.Module,name: str = '',depth_first: bool = True,include_root: bool = False) -> nn.Moduleworldfoundry.core.named_applyfrom worldfoundry.core import named_apply简介
named_apply — Apply `fn(module=..., name=...) recursively with stable module names. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：nn.Module`。
源码 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.
参数
fnCallable[..., object]modulenn.Modulenamestr- 默认值:
'' depth_firstbool- 默认值:
True include_rootbool- 默认值:
False
返回值: nn.Module
PatchEmbed
clsclass 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)worldfoundry.core.PatchEmbedfrom worldfoundry.core import PatchEmbed简介
PatchEmbed — 2D image patch embedding: `(B, C, H, W) -> (B, N, D)`. 属于 Core 神经网络/数学辅助（共享张量变换）。
参数
img_sizeUnion[int, tuple[int, int]]- Nominal image height/width used to report patch count.默认值:
224 patch_sizeUnion[int, tuple[int, int]]- Convolution kernel/stride height and width.默认值:
16 in_chansint- Input channel count.默认值:
3 embed_dimint- Output token width.默认值:
768 norm_layerOptional[Callable[..., nn.Module]]- Optional normalization factory applied per token.默认值:
None flatten_embeddingbool- Return `
(B, N, D)when true, otherwise(B, grid_h, grid_w, D)`.默认值:True
方法
简介
forward — Project a divisible NCHW image batch into patch embeddings.
参数
xTensor
异常
ValueError- Runtime height or width is not divisible by patch size.
返回值: Tensor
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)worldfoundry.core.PatchEmbed_Mlpfrom worldfoundry.core import PatchEmbed_Mlp简介
PatchEmbed_Mlp — Patch embedding implemented with pixel unshuffle and MLP projection. 属于 Core 神经网络/数学辅助（共享张量变换）。
参数
img_sizeUnion[int, tuple[int, int]]- 默认值:
224 patch_sizeUnion[int, tuple[int, int]]- 默认值:
16 in_chansint- 默认值:
3 embed_dimint- 默认值:
768 norm_layerOptional[Callable[..., nn.Module]]- 默认值:
None flatten_embeddingbool- 默认值:
True
class PatchGridSpec(original_shape: tuple[int, ...],patch_size: tuple[int, int],layout: ImageLayout = 'nchw')worldfoundry.core.PatchGridSpecfrom worldfoundry.core import PatchGridSpec简介
PatchGridSpec — Shape contract needed to invert a 2D image patchification. 属于 Core 神经网络/数学辅助（共享张量变换）。
属性
original_shapetuple[int, ...]patch_sizetuple[int, int]layoutImageLayout- 默认值:
'nchw'
方法
patchify_image
funcdef patchify_image(value: Any,patch_size: int | Sequence[int],layout: ImageLayout = 'nchw') -> tuple[Any, PatchGridSpec]worldfoundry.core.patchify_imagefrom worldfoundry.core import patchify_image简介
patchify_image — Convert an image tensor into flattened 2D patch tokens. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：tuple[Any, PatchGridSpec]。
源码 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)`.
参数
valueAnypatch_sizeint | Sequence[int]layoutImageLayout- 默认值:
'nchw'
返回值: tuple[Any, PatchGridSpec]
class PixelUnshuffle(downscale_factor: int)worldfoundry.core.PixelUnshufflefrom worldfoundry.core import PixelUnshuffle简介
PixelUnshuffle — Module wrapper for `torch.nn.functional.pixel_unshuffle`. 属于 Core 神经网络/数学辅助（共享张量变换）。
参数
downscale_factorint
方法
def quaternion_xyzw_to_rotation_matrix(quaternions: Tensor) -> Tensorworldfoundry.core.quaternion_xyzw_to_rotation_matrixfrom worldfoundry.core import quaternion_xyzw_to_rotation_matrix简介
quaternion_xyzw_to_rotation_matrix — Convert `[..., x, y, z, w] quaternions to rotation matrices. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：Tensor`。
参数
quaternionsTensor
返回值: Tensor
ray_condition
funcdef ray_condition(K,c2w,H: int,W: int,device,flip_flag = None,use_ray_o: bool = False)worldfoundry.core.ray_conditionfrom worldfoundry.core import ray_condition简介
ray_condition — Build per-pixel Plucker ray features from intrinsics and camera-to-world poses. 属于 Core 神经网络/数学辅助（共享张量变换）。
参数
K- Camera intrinsics in `
[fx, fy, cx, cy]layout with shape[B,V,4]. WhenNone`, rays use a constant forward camera-space direction. c2w- Camera-to-world matrices with shape `
[B,V,4,4]`. Hint- Output grid height.
Wint- Output grid width.
device- Device for generated coordinate grids.
flip_flag- Optional boolean mask selecting horizontally flipped views.默认值:
None use_ray_obool- If true, concatenate `
[ray_origin, ray_direction]instead of the Plucker[ray_direction x ray_origin, ray_direction]` form.默认值:False
rms_norm
funcdef rms_norm(value: Any,weight: Any = None,eps: float = 1e-06) -> Anyworldfoundry.core.rms_normfrom worldfoundry.core import rms_norm简介
rms_norm — Apply RMS normalization over the last dimension. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：Any。
参数
valueAnyweightAny- 默认值:
None epsfloat- 默认值:
1e-06
返回值: Any
def rotation_matrix_to_euler_angles_opencv(rotation: np.ndarray) -> tuple[float, float, float]worldfoundry.core.rotation_matrix_to_euler_angles_opencvfrom worldfoundry.core import rotation_matrix_to_euler_angles_opencv简介
rotation_matrix_to_euler_angles_opencv — Return the ZYX angles in OpenCV's `(z, y, x) order. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：tuple[float, float, float]`。
参数
rotationnp.ndarray
返回值: tuple[float, float, float]
def rotation_matrix_to_euler_angles_zyx(rotation: np.ndarray) -> np.ndarrayworldfoundry.core.rotation_matrix_to_euler_angles_zyxfrom worldfoundry.core import rotation_matrix_to_euler_angles_zyx简介
rotation_matrix_to_euler_angles_zyx — Convert a 3x3 rotation matrix to `[x, y, z] Euler angles. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：np.ndarray`。
参数
rotationnp.ndarray
返回值: np.ndarray
def rotation_matrix_to_quaternion_wxyz(rotation: np.ndarray) -> np.ndarrayworldfoundry.core.rotation_matrix_to_quaternion_wxyzfrom worldfoundry.core import rotation_matrix_to_quaternion_wxyz简介
rotation_matrix_to_quaternion_wxyz — Convert a 3x3 rotation matrix to a normalized `[w, x, y, z] quaternion. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：np.ndarray`。
参数
rotationnp.ndarray
返回值: np.ndarray
def rotation_matrix_to_quaternion_xyzw(matrix: Tensor) -> Tensorworldfoundry.core.rotation_matrix_to_quaternion_xyzwfrom worldfoundry.core import rotation_matrix_to_quaternion_xyzw简介
rotation_matrix_to_quaternion_xyzw — Convert `[..., 3, 3] rotation matrices to canonical XYZW quaternions. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：Tensor`。
源码 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.
参数
matrixTensor
返回值: Tensor
class SchedulerInterface(alphas_cumprod: torch.Tensor)worldfoundry.core.SchedulerInterfacefrom worldfoundry.core import SchedulerInterface简介
SchedulerInterface — Base interface for diffusion noise schedules. 属于 Core 神经网络/数学辅助（共享张量变换）。
属性
alphas_cumprodtorch.Tensor
方法
简介
add_noise — Run the forward corruption process.
参数
clean_latenttorch.Tensornoisetorch.Tensortimesteptorch.Tensor
简介
convert_x0_to_noise — Convert a clean-data prediction to a noise prediction.
参数
x0torch.Tensorxttorch.Tensortimesteptorch.Tensor
返回值: torch.Tensor
convert_noise_to_x0(noise: torch.Tensor,xt: torch.Tensor,timestep: torch.Tensor) -> torch.Tensor源码简介
convert_noise_to_x0 — Convert a noise prediction to a clean-data prediction.
参数
noisetorch.Tensorxttorch.Tensortimesteptorch.Tensor
返回值: torch.Tensor
convert_velocity_to_x0(velocity: torch.Tensor,xt: torch.Tensor,timestep: torch.Tensor) -> torch.Tensor源码简介
convert_velocity_to_x0 — Convert a velocity prediction to a clean-data prediction.
参数
velocitytorch.Tensorxttorch.Tensortimesteptorch.Tensor
返回值: torch.Tensor
def split_attention_heads(value: Any, num_heads: int) -> Anyworldfoundry.core.split_attention_headsfrom worldfoundry.core import split_attention_heads简介
split_attention_heads — Reshape `(..., seq, hidden) to (..., heads, seq, head_dim). 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：Any`。
参数
valueAnynum_headsint
返回值: Any
def standardize_quaternion_xyzw(quaternions: Tensor) -> Tensorworldfoundry.core.standardize_quaternion_xyzwfrom worldfoundry.core import standardize_quaternion_xyzw简介
standardize_quaternion_xyzw — Choose the equivalent XYZW quaternion whose real part is non-negative. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：Tensor。
参数
quaternionsTensor
返回值: Tensor
SwiGLUFFN
clsclass 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)worldfoundry.core.SwiGLUFFNfrom worldfoundry.core import SwiGLUFFN简介
各模型接入可复用的 SwiGLU 前馈层。
参数
in_featuresintout_featuresOptional[int]- 默认值:
None act_layerCallable[..., nn.Module] | None- 默认值:
None dropfloat- 默认值:
0.0 biasbool- 默认值:
True
方法
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)worldfoundry.core.SwiGLUFFNFusedfrom worldfoundry.core import SwiGLUFFNFused简介
SwiGLUFFNFused — SwiGLU FFN with hidden width rounded for tensor-core-friendly matmuls. 属于 Core 神经网络/数学辅助（共享张量变换）。
参数
in_featuresint- Input width.
out_featuresOptional[int]- Output width; defaults to `
in_features`.默认值:None act_layerCallable[..., nn.Module] | None- Accepted for MLP-constructor compatibility; unused.默认值:
None dropfloat- Accepted for MLP-constructor compatibility; unused.默认值:
0.0 biasbool- Enable biases in the fused input and output projections.默认值:
True
def transformer_shape_spec(hidden_size: int,num_heads: int,mlp_ratio: float = 4.0,multiple_of: int | None = None) -> TransformerShapeSpecworldfoundry.core.transformer_shape_specfrom worldfoundry.core import transformer_shape_spec简介
transformer_shape_spec — Build a reusable shape spec for attention and MLP dimensions. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：TransformerShapeSpec。
参数
num_headsintmlp_ratiofloat- 默认值:
4.0 multiple_ofint | None- 默认值:
None
返回值: TransformerShapeSpec
class TransformerShapeSpec(hidden_size: int,num_heads: int,head_dim: int,mlp_hidden_size: int)worldfoundry.core.TransformerShapeSpecfrom worldfoundry.core import TransformerShapeSpec简介
TransformerShapeSpec — Common transformer dimensions derived from hidden size and head count. 属于 Core 神经网络/数学辅助（共享张量变换）。
属性
num_headsinthead_dimint
unpatchify_image
funcdef unpatchify_image(patches: Any, spec: PatchGridSpec) -> Anyworldfoundry.core.unpatchify_imagefrom worldfoundry.core import unpatchify_image简介
unpatchify_image — Invert `patchify_image using the returned PatchGridSpec. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：Any`。
参数
patchesAnyspecPatchGridSpec
返回值: Any
zero_module
funcdef zero_module(module: nn.Module) -> nn.Moduleworldfoundry.core.zero_modulefrom worldfoundry.core import zero_module简介
zero_module — Detach and zero all parameters in a module. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：nn.Module。
参数
modulenn.Module
返回值: nn.Module