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。
完整参考
以下为该类别的生成签名。可用本页符号索引跳转；源码链接指向各惰性导出背后的具体实现。
38 个公开符号
def attention_head_dim(hidden_size: int,num_heads: int) -> intattention_head_dim — Return per-head attention width and validate divisibility. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：int。
参数
num_headsint返回值
def causal_attention_mask(query_len: int,key_len: int | None = None,include_self: bool = True) -> np.ndarray简介
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默认值 Noneinclude_selfbool默认值 True返回值
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.Tensor简介
clip_grad_norm_ — Clip gradients and optionally reduce the norm across pipeline-parallel stages. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：torch.Tensor。
源码 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.
参数
parameterstorch.Tensor | Iterable[torch.Tensor]max_normfloatnorm_typefloat默认值 2.0error_if_nonfinitebool默认值 Falseforeachbool | None默认值 Nonepp_meshtorch.distributed.device_mesh.DeviceMesh | None默认值 None返回值
def clip_grads_with_norm_(parameters: torch.Tensor | Iterable[torch.Tensor],max_norm: float,total_norm: torch.Tensor,foreach: bool | None = None) -> None简介
clip_grads_with_norm_ — Scale parameter gradients in-place using a precomputed total norm. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：None。
源码 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.
参数
parameterstorch.Tensor | Iterable[torch.Tensor]max_normfloattotal_normtorch.Tensorforeachbool | None默认值 None返回值
class DropPath(drop_prob: float | None = 0.0,scale_by_keep: bool = True)训练时按样本丢弃残差路径的 Stochastic Depth 模块。
参数
drop_probfloat | None默认值 0.0scale_by_keepbool默认值 True方法
forward — No-op in eval; in train, scale surviving samples when requested.
参数
xTensor返回值
def euler_angles_to_rotation_matrix_zyx(euler_angles: np.ndarray) -> np.ndarrayeuler_angles_to_rotation_matrix_zyx — Convert `[x, y, z] Euler angles to Rz @ Ry @ Rx. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：np.ndarray`。
参数
euler_anglesnp.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,exponential_shift: bool = False,exponential_shift_mu: float | None = None,shift_terminal: float | None = None)简介
FlowMatchScheduler — Flow-matching scheduler shared by bundled video runtimes. 属于 Core 神经网络/数学辅助（共享张量变换）。
源码 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.
参数
num_inference_stepsint默认值 100num_train_timestepsint默认值 1000shiftfloat默认值 3.0sigma_maxfloat默认值 1.0sigma_minfloat默认值 0.003 / 1.002inverse_timestepsbool默认值 Falseextra_one_stepbool默认值 Falsereverse_sigmasbool默认值 Falseexponential_shiftbool默认值 Falseexponential_shift_mufloat | None默认值 Noneshift_terminalfloat | None默认值 None方法
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) -> None源码set_timesteps — Rebuild `sigmas / timesteps` for a new step count or shift.
参数
num_inference_stepsint默认值 100denoising_strengthfloat默认值 1.0trainingbool默认值 Falseshiftfloat | None默认值 Nonedynamic_shift_lenint | None默认值 Noneexponential_shift_mufloat | None默认值 None返回值
简介
step — One Euler step: `sample + velocity * (σ_next - σ)`.
源码 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.
参数
model_outputtimestepsampleto_finalbool默认值 Falsereturn_to_timestep — Recover a velocity that maps `sample to stabilized_sample`.
参数
timestepsamplestabilized_sampleadd_noise — Run the forward corruption process.
参数
original_samplesnoisetimesteptraining_target — Flow-matching target `noise - sample` (velocity toward noise).
参数
samplenoisetimesteptraining_weight — Per-timestep loss weight from the Gaussian bump built in training mode.
参数
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) -> float源码calculate_shift — Interpolate exponential-shift μ from sequence length (Flux-style).
参数
image_seq_lenintbase_seq_lenint默认值 256max_seq_lenint默认值 8192base_shiftfloat默认值 0.5max_shiftfloat默认值 0.9返回值
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.Tensor简介
get_total_norm — Compute the total norm of tensors as if their flattened values were concatenated. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：torch.Tensor。
源码 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.
参数
tensorstorch.Tensor | Iterable[torch.Tensor]norm_typefloat默认值 2.0error_if_nonfinitebool默认值 Falseforeachbool | None默认值 None返回值
def layer_scale(value: Any,scale: Any,bias: Any = None) -> Anylayer_scale — Apply a stateless elementwise layer-scale transform. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：Any。
参数
valueAnyscaleAnybiasAny默认值 None返回值
class LayerScale(dim: int,init_values: Union[float, Tensor] = 1e-05,inplace: bool = False,device = None)现代 ViT block 中常见的可学习逐通道残差缩放。
参数
dimintinit_valuesUnion[float, Tensor]默认值 1e-05gamma`.inplacebool默认值 Falseforward`.device默认值 None方法
reset_parameters — Copy a tensor init or fill a scalar; used after meta-device materialize.
返回值
forward — Scale the final dimension of `x by the learned gamma`.
参数
xTensor返回值
extra_repr — Include init and inplace flags so module printouts distinguish LayerScales.
返回值
def merge_attention_heads(value: Any) -> Anymerge_attention_heads — Invert `split_attention_heads for (..., heads, seq, head_dim) tensors. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：Any`。
参数
valueAny返回值
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)Mlp — timm-style 2-layer ViT FFN (fc1 → act → fc2); distinct from `SamHeadMLP`. 属于 Core 神经网络/数学辅助（共享张量变换）。
参数
in_featuresintout_featuresOptional[int]默认值 Nonein_features`.act_layerCallable[..., nn.Module]默认值 nn.GELUdropfloat | tuple[float, float]默认值 0.0biasbool | tuple[bool, bool]默认值 Truedevice默认值 None方法
forward — Apply `fc1 → activation → dropout → fc2 → dropout`.
参数
xTensor返回值
def named_apply(fn: Callable[..., object],module: nn.Module,name: str = '',depth_first: bool = True,include_root: bool = False) -> nn.Module简介
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默认值 Trueinclude_rootbool默认值 False返回值
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)简介
named_camera_trajectory_tensors — Build the eight named GEN3C camera trajectories without a model runtime. 属于 Core 神经网络/数学辅助（共享张量变换）。
源码 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.
参数
trajectorystrinitial_world_to_cameraAnyinitial_intrinsicAnynum_framesintmovement_distancefloat默认值 0.3camera_rotationstr默认值 'center_facing'center_depthfloat默认值 1.0num_circlesint默认值 1class 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)PatchEmbed — 2D image patch embedding: `(B, C, H, W) -> (B, N, D)`. 属于 Core 神经网络/数学辅助（共享张量变换）。
参数
img_sizeUnion[int, tuple[int, int]]默认值 224patch_sizeUnion[int, tuple[int, int]]默认值 16in_chansint默认值 3embed_dimint默认值 768norm_layerOptional[Callable[..., nn.Module]]默认值 Noneflatten_embeddingbool默认值 True(B, N, D) when true, otherwise (B, grid_h, grid_w, D)`.方法
reset_parameters — Uniform init scaled by fan-in of the patch convolution (checkpoint match).
返回值
forward — Project a divisible NCHW image batch into patch embeddings.
参数
xTensor返回值
异常
ValueError- Runtime height or width is not divisible by patch size.
flops — Report conv + optional norm FLOPs for the *nominal* `img_size` grid.
返回值
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)PatchEmbed_Mlp — Patch embedding implemented with pixel unshuffle and MLP projection. 属于 Core 神经网络/数学辅助（共享张量变换）。
参数
img_sizeUnion[int, tuple[int, int]]默认值 224patch_sizeUnion[int, tuple[int, int]]默认值 16in_chansint默认值 3embed_dimint默认值 768norm_layerOptional[Callable[..., nn.Module]]默认值 Noneflatten_embeddingbool默认值 Trueclass PatchGridSpec(original_shape: tuple[int, ...],patch_size: tuple[int, int],layout: ImageLayout = 'nchw')PatchGridSpec — Shape contract needed to invert a 2D image patchification. 属于 Core 神经网络/数学辅助（共享张量变换）。
属性
original_shapetuple[int, ...]patch_sizetuple[int, int]layoutImageLayout默认值 'nchw'方法
batch_shape — Leading axes before the CHW / HWC triple (may be empty).
参数
self返回值
spatial_shape — `(H, W)` regardless of NCHW vs NHWC.
参数
self返回值
grid_shape — Number of patches along height and width.
参数
self返回值
patch_vector_size — Flattened per-patch width `C * ph * pw`.
参数
self返回值
patch_count — `grid_h * grid_w` — the sequence length after patchify.
参数
self返回值
def patchify_image(value: Any,patch_size: int | Sequence[int],layout: ImageLayout = 'nchw') -> tuple[Any, PatchGridSpec]简介
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'返回值
class PixelUnshuffle(downscale_factor: int)PixelUnshuffle — Module wrapper for `torch.nn.functional.pixel_unshuffle`. 属于 Core 神经网络/数学辅助（共享张量变换）。
参数
downscale_factorint方法
forward — Unshuffle `[B, C, H, W]; empty tensors use view` because F.
参数
valueTensor返回值
异常
ValueError- empty input spatial dims are 0 or not divisible by the factor.
def quaternion_xyzw_to_rotation_matrix(quaternions: Tensor) -> Tensorquaternion_xyzw_to_rotation_matrix — Convert `[..., x, y, z, w] quaternions to rotation matrices. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：Tensor`。
参数
quaternionsTensor返回值
def ray_condition(K,c2w,H: int,W: int,device,flip_flag = None,use_ray_o: bool = False)ray_condition — Build per-pixel Plucker ray features from intrinsics and camera-to-world poses. 属于 Core 神经网络/数学辅助（共享张量变换）。
参数
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_flag默认值 Noneuse_ray_obool默认值 False[ray_origin, ray_direction] instead of the Plucker [ray_origin x ray_direction, ray_direction]` form.def rotation_matrix_to_euler_angles_opencv(rotation: np.ndarray) -> tuple[float, float, float]rotation_matrix_to_euler_angles_opencv — Return the ZYX angles in OpenCV's `(z, y, x) order. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：tuple[float, float, float]`。
参数
rotationnp.ndarray返回值
def rotation_matrix_to_euler_angles_zyx(rotation: np.ndarray) -> np.ndarrayrotation_matrix_to_euler_angles_zyx — Convert a 3x3 rotation matrix to `[x, y, z] Euler angles. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：np.ndarray`。
参数
rotationnp.ndarray返回值
def rotation_matrix_to_quaternion_wxyz(rotation: np.ndarray) -> np.ndarrayrotation_matrix_to_quaternion_wxyz — Convert a 3x3 rotation matrix to a normalized `[w, x, y, z] quaternion. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：np.ndarray`。
参数
rotationnp.ndarray返回值
def rotation_matrix_to_quaternion_xyzw(matrix: Tensor) -> Tensor简介
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返回值
class SchedulerInterface(alphas_cumprod: torch.Tensor)SchedulerInterface — Base interface for diffusion noise schedules. 属于 Core 神经网络/数学辅助（共享张量变换）。
属性
alphas_cumprodtorch.Tensor方法
add_noise — Run the forward corruption process.
参数
clean_latenttorch.Tensornoisetorch.Tensortimesteptorch.Tensorconvert_x0_to_noise — Convert a clean-data prediction to a noise prediction.
参数
x0torch.Tensorxttorch.Tensortimesteptorch.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返回值
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返回值
def split_attention_heads(value: Any,num_heads: int) -> Anysplit_attention_heads — Reshape `(..., seq, hidden) to (..., heads, seq, head_dim). 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：Any`。
参数
valueAnynum_headsint返回值
def standardize_quaternion_xyzw(quaternions: Tensor) -> Tensorstandardize_quaternion_xyzw — Choose the equivalent XYZW quaternion whose real part is non-negative. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：Tensor。
参数
quaternionsTensor返回值
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 前馈层。
参数
in_featuresintout_featuresOptional[int]默认值 Noneact_layerCallable[..., nn.Module] | None默认值 Nonedropfloat默认值 0.0biasbool默认值 True方法
forward — Last-dim SwiGLU matching the `w12 / w3` checkpoint split.
参数
xTensor返回值
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)SwiGLUFFNFused — SwiGLU FFN with hidden width rounded for tensor-core-friendly matmuls. 属于 Core 神经网络/数学辅助（共享张量变换）。
参数
in_featuresintout_featuresOptional[int]默认值 Nonein_features`.act_layerCallable[..., nn.Module] | None默认值 Nonedropfloat默认值 0.0biasbool默认值 Truedef transformer_shape_spec(hidden_size: int,num_heads: int,mlp_ratio: float = 4.0,multiple_of: int | None = None) -> TransformerShapeSpectransformer_shape_spec — Build a reusable shape spec for attention and MLP dimensions. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：TransformerShapeSpec。
参数
num_headsintmlp_ratiofloat默认值 4.0multiple_ofint | None默认值 Noneclass TransformerShapeSpec(hidden_size: int,num_heads: int,head_dim: int,mlp_hidden_size: int)TransformerShapeSpec — Common transformer dimensions derived from hidden size and head count. 属于 Core 神经网络/数学辅助（共享张量变换）。
属性
num_headsinthead_dimintdef unpatchify_image(patches: Any,spec: PatchGridSpec) -> Anyunpatchify_image — Invert `patchify_image using the returned PatchGridSpec. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：Any`。
参数
patchesAnyspecPatchGridSpec返回值
def zero_module(module: nn.Module) -> nn.Modulezero_module — Detach and zero all parameters in a module. 属于 Core 神经网络/数学辅助（共享张量变换）。 标注返回类型：nn.Module。
参数
modulenn.Module