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.
37 public symbols
def attention_head_dim(hidden_size: int, num_heads: int) -> intworldfoundry.core.attention_head_dimfrom worldfoundry.core import attention_head_dimOverview
Return per-head attention width and validate divisibility. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: int.
Parameters
num_headsint
Returns: 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_maskOverview
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 | None- default:
None include_selfbool- default:
True
Returns: 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_Overview
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.
Parameters
parameterstorch.Tensor | Iterable[torch.Tensor]max_normfloatnorm_typefloat- default:
2.0 error_if_nonfinitebool- default:
False foreachbool | None- default:
None pp_meshtorch.distributed.device_mesh.DeviceMesh | None- default:
None
Returns: 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_Overview
Scale parameter gradients in-place using a precomputed total norm. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: None.
Parameters
parameterstorch.Tensor | Iterable[torch.Tensor]max_normfloattotal_normtorch.Tensorforeachbool | None- default:
None
Returns: 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_pathOverview
Functional form of per-sample stochastic depth for residual branches.
Parameters
xTensordrop_probfloat- default:
0.0 trainingbool- default:
False scale_by_keepbool- default:
True
Returns: Tensor
DropPath
clsclass DropPath(drop_prob: float | None = 0.0, scale_by_keep: bool = True)worldfoundry.core.DropPathfrom worldfoundry.core import DropPathOverview
Stochastic depth module that drops residual paths per sample during training.
Parameters
drop_probfloat | None- default:
0.0 scale_by_keepbool- default:
True
Methods
Overview
Public method on this type.
Parameters
xTensor
Returns: 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_zyxOverview
Convert `[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.ndarray
Returns: 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 FlowMatchSchedulerOverview
Flow-matching scheduler shared by bundled video runtimes. Belongs to Core neural-net/math helpers (shared tensor transforms).
Parameters
num_inference_stepsint- default:
100 num_train_timestepsint- default:
1000 shiftfloat- default:
3.0 sigma_maxfloat- default:
1.0 sigma_minfloat- default:
0.003 / 1.002 inverse_timestepsbool- default:
False extra_one_stepbool- default:
False reverse_sigmasbool- default:
False
Methods
set_timesteps(num_inference_steps: int = 100,denoising_strength: float = 1.0,training: bool = False) -> NonesourceOverview
Public method on this type.
Parameters
num_inference_stepsint- default:
100 denoising_strengthfloat- default:
1.0 trainingbool- default:
False
Returns: None
Overview
Public method on this type.
Parameters
model_outputtimestepsampleto_finalbool- default:
False
Overview
Run the forward corruption process.
Parameters
original_samplesnoisetimestep
Overview
Public method on this type.
Parameters
samplenoisetimestep
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_normOverview
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.
Parameters
tensorstorch.Tensor | Iterable[torch.Tensor]norm_typefloat- default:
2.0 error_if_nonfinitebool- default:
False foreachbool | None- default:
None
Returns: torch.Tensor
layer_scale
funcdef layer_scale(value: Any,scale: Any,bias: Any = None) -> Anyworldfoundry.core.layer_scalefrom worldfoundry.core import layer_scaleOverview
Apply a stateless elementwise layer-scale transform. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: Any.
Parameters
valueAnyscaleAnybiasAny- default:
None
Returns: Any
LayerScale
clsclass LayerScale(dim: int,init_values: Union[float, Tensor] = 1e-05,inplace: bool = False,device = None)worldfoundry.core.LayerScalefrom worldfoundry.core import LayerScaleOverview
Learnable per-channel residual scaling used in modern Vision Transformer blocks.
Parameters
dimint- Channel width of the final tensor dimension.
init_valuesUnion[float, Tensor]- Scalar or per-channel initialization for `
gamma`.default:1e-05 inplacebool- Multiply the input in place during `
forward`.default:False device- Optional parameter device.default:
None
def merge_attention_heads(value: Any) -> Anyworldfoundry.core.merge_attention_headsfrom worldfoundry.core import merge_attention_headsOverview
Invert `split_attention_heads for (..., heads, seq, head_dim) tensors. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: Any`.
Parameters
valueAny
Returns: 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 MlpOverview
timm-style 2-layer ViT FFN (fc1 β act β fc2); distinct from `SamHeadMLP`. Belongs to Core neural-net/math helpers (shared tensor transforms).
Parameters
in_featuresint- Input width.
out_featuresOptional[int]- Output width; defaults to `
in_features`.default:None act_layerCallable[..., nn.Module]- Activation module factory between projections.default:
nn.GELU dropfloat | tuple[float, float]- One probability for both dropout sites or a pair for the first and second sites.default:
0.0 biasbool | tuple[bool, bool]- One bias flag for both projections or a pair.default:
True device- Optional projection parameter device.default:
None
Methods
Overview
Apply `fc1 β activation β dropout β fc2 β dropout`.
Parameters
xTensor
Returns: 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_applyOverview
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.Modulenamestr- default:
'' depth_firstbool- default:
True include_rootbool- default:
False
Returns: 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 PatchEmbedOverview
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]]- Nominal image height/width used to report patch count.default:
224 patch_sizeUnion[int, tuple[int, int]]- Convolution kernel/stride height and width.default:
16 in_chansint- Input channel count.default:
3 embed_dimint- Output token width.default:
768 norm_layerOptional[Callable[..., nn.Module]]- Optional normalization factory applied per token.default:
None flatten_embeddingbool- Return `
(B, N, D)when true, otherwise(B, grid_h, grid_w, D)`.default:True
Methods
Overview
Project a divisible NCHW image batch into patch embeddings.
Parameters
xTensor
Raises
ValueError- Runtime height or width is not divisible by patch size.
Returns: 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_MlpOverview
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:
224 patch_sizeUnion[int, tuple[int, int]]- default:
16 in_chansint- default:
3 embed_dimint- default:
768 norm_layerOptional[Callable[..., nn.Module]]- default:
None flatten_embeddingbool- default:
True
class PatchGridSpec(original_shape: tuple[int, ...],patch_size: tuple[int, int],layout: ImageLayout = 'nchw')worldfoundry.core.PatchGridSpecfrom worldfoundry.core import PatchGridSpecOverview
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]layoutImageLayout- default:
'nchw'
Methods
Overview
Public property on this type.
Parameters
self
Returns: tuple[int, ...]
Overview
Public property on this type.
Parameters
self
Returns: tuple[int, int]
Overview
Public property on this type.
Parameters
self
Returns: tuple[int, int]
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_imageOverview
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]layoutImageLayout- default:
'nchw'
Returns: tuple[Any, PatchGridSpec]
Permute
clsclass Permute(dims: tuple[int, ])worldfoundry.core.Permutefrom worldfoundry.core import PermuteOverview
Module wrapper around `Tensor.permute`. Belongs to Core neural-net/math helpers (shared tensor transforms).
Attributes
dimstuple[int, ...]
Methods
Overview
Public method on this type.
Parameters
valueTensor
Returns: Tensor
class PixelUnshuffle(downscale_factor: int)worldfoundry.core.PixelUnshufflefrom worldfoundry.core import PixelUnshuffleOverview
Module wrapper for `torch.nn.functional.pixel_unshuffle`. Belongs to Core neural-net/math helpers (shared tensor transforms).
Parameters
downscale_factorint
Methods
Overview
Public method on this type.
Parameters
valueTensor
Returns: Tensor
def quaternion_xyzw_to_rotation_matrix(quaternions: Tensor) -> Tensorworldfoundry.core.quaternion_xyzw_to_rotation_matrixfrom worldfoundry.core import quaternion_xyzw_to_rotation_matrixOverview
Convert `[..., x, y, z, w] quaternions to rotation matrices. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: Tensor`.
Parameters
quaternionsTensor
Returns: 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_conditionOverview
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- 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.default:
None use_ray_obool- If true, concatenate `
[ray_origin, ray_direction]instead of the Plucker[ray_direction x ray_origin, ray_direction]` form.default:False
rms_norm
funcdef rms_norm(value: Any,weight: Any = None,eps: float = 1e-06) -> Anyworldfoundry.core.rms_normfrom worldfoundry.core import rms_normOverview
Apply RMS normalization over the last dimension. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: Any.
Parameters
valueAnyweightAny- default:
None epsfloat- default:
1e-06
Returns: 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_opencvOverview
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.ndarray
Returns: 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_zyxOverview
Convert 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.ndarray
Returns: 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_wxyzOverview
Convert 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.ndarray
Returns: np.ndarray
def rotation_matrix_to_quaternion_xyzw(matrix: Tensor) -> Tensorworldfoundry.core.rotation_matrix_to_quaternion_xyzwfrom worldfoundry.core import rotation_matrix_to_quaternion_xyzwOverview
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
matrixTensor
Returns: Tensor
class SchedulerInterface(alphas_cumprod: torch.Tensor)worldfoundry.core.SchedulerInterfacefrom worldfoundry.core import SchedulerInterfaceOverview
Base interface for diffusion noise schedules. Belongs to Core neural-net/math helpers (shared tensor transforms).
Attributes
alphas_cumprodtorch.Tensor
Methods
Overview
Run the forward corruption process.
Parameters
clean_latenttorch.Tensornoisetorch.Tensortimesteptorch.Tensor
convert_x0_to_noise(x0: torch.Tensor,xt: torch.Tensor,timestep: torch.Tensor) -> torch.TensorsourceOverview
Convert a clean-data prediction to a noise prediction.
Parameters
x0torch.Tensorxttorch.Tensortimesteptorch.Tensor
Returns: torch.Tensor
convert_noise_to_x0(noise: torch.Tensor,xt: torch.Tensor,timestep: torch.Tensor) -> torch.TensorsourceOverview
Convert a noise prediction to a clean-data prediction.
Parameters
noisetorch.Tensorxttorch.Tensortimesteptorch.Tensor
Returns: torch.Tensor
convert_velocity_to_x0(velocity: torch.Tensor,xt: torch.Tensor,timestep: torch.Tensor) -> torch.TensorsourceOverview
Convert a velocity prediction to a clean-data prediction.
Parameters
velocitytorch.Tensorxttorch.Tensortimesteptorch.Tensor
Returns: torch.Tensor
def split_attention_heads(value: Any, num_heads: int) -> Anyworldfoundry.core.split_attention_headsfrom worldfoundry.core import split_attention_headsOverview
Reshape `(..., seq, hidden) to (..., heads, seq, head_dim). Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: Any`.
Parameters
valueAnynum_headsint
Returns: Any
def standardize_quaternion_xyzw(quaternions: Tensor) -> Tensorworldfoundry.core.standardize_quaternion_xyzwfrom worldfoundry.core import standardize_quaternion_xyzwOverview
Choose 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
quaternionsTensor
Returns: 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 SwiGLUFFNOverview
SwiGLU feed-forward layer shared across model integrations.
Parameters
in_featuresintout_featuresOptional[int]- default:
None act_layerCallable[..., nn.Module] | None- default:
None dropfloat- default:
0.0 biasbool- default:
True
Methods
Overview
Public method on this type.
Parameters
xTensor
Returns: Tensor
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 SwiGLUFFNFusedOverview
SwiGLU FFN with hidden width rounded for tensor-core-friendly matmuls. Belongs to Core neural-net/math helpers (shared tensor transforms).
Parameters
in_featuresint- Input width.
out_featuresOptional[int]- Output width; defaults to `
in_features`.default:None act_layerCallable[..., nn.Module] | None- Accepted for MLP-constructor compatibility; unused.default:
None dropfloat- Accepted for MLP-constructor compatibility; unused.default:
0.0 biasbool- Enable biases in the fused input and output projections.default:
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_specOverview
Build 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_ratiofloat- default:
4.0 multiple_ofint | None- default:
None
Returns: TransformerShapeSpec
class TransformerShapeSpec(hidden_size: int,num_heads: int,head_dim: int,mlp_hidden_size: int)worldfoundry.core.TransformerShapeSpecfrom worldfoundry.core import TransformerShapeSpecOverview
Common transformer dimensions derived from hidden size and head count. Belongs to Core neural-net/math helpers (shared tensor transforms).
Attributes
num_headsinthead_dimint
unpatchify_image
funcdef unpatchify_image(patches: Any, spec: PatchGridSpec) -> Anyworldfoundry.core.unpatchify_imagefrom worldfoundry.core import unpatchify_imageOverview
Invert `patchify_image using the returned PatchGridSpec. Belongs to Core neural-net/math helpers (shared tensor transforms). Annotated return type: Any`.
Parameters
patchesAnyspecPatchGridSpec
Returns: Any
zero_module
funcdef zero_module(module: nn.Module) -> nn.Moduleworldfoundry.core.zero_modulefrom worldfoundry.core import zero_moduleOverview
Detach 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
Returns: nn.Module