nn

Module for PyTorch neural network utilities.

Classes:

Name	Description
ModuleDefaultDict	Holds submodules in a dictionary and calls a factory function to supply missing values.
TiedLinear	A bias-free linear layer whose weight is shared with another module, read live on every call.

ModuleDefaultDict

Bases: ModuleDict, Generic[ModuleT]

Holds submodules in a dictionary and calls a factory function to supply missing values.

Parameters:

Name	Type	Description	Default
default_factory	Callable[[], ModuleT] | None	A no argument constructor for missing Module items. If None, missing items will raise a KeyError.	None
**kwargs	ModuleT	Initial dictionary values.	{}

Methods:

Name	Description
__getitem__	Retrieve a module by key, creating it if missing.

__getitem__

__getitem__(key: str) -> ModuleT

Retrieve a module by key, creating it if missing.

Parameters:

Name	Type	Description	Default
key	str	The name of the module to retrieve.	required

Returns:

Type	Description
ModuleT	The requested module.

TiedLinear

A bias-free linear layer whose weight is shared with another module, read live on every call.

TiedLinear does not copy or cache tied_module's weight; it reads tied_module.weight on every call. This keeps the two layers tied even if tied_module is later re-wrapped by something like FSDP, which replaces the memory backing a weight but not the Module object that owns it.

TiedLinear intentionally does not subclass torch.nn.Module. Module.__setattr__ auto-registers any Module-valued attribute as a submodule, which would cause tied_module to be registered twice (once under its original owner, once under TiedLinear), duplicating its weight in state_dict() and parameters(). The actual torch.nn.functional.linear call is performed by the linear attribute, a real Module, so forward hooks can still be attached to observe the transform.

Parameters:

Name	Type	Description	Default
tied_module	Module	The module whose weight attribute is shared. Only the weight is used; any bias on tied_module is ignored.	required

Raises:

Type	Description
AttributeError	If tied_module does not have a weight attribute.

Added in version v3.58.0. Native replacement for torchtune's `TiedLinear`, removing the `torchtune` dependency.

Methods:

Name	Description
__call__	Apply the tied linear transform.

__call__

__call__(x: Tensor) -> torch.Tensor

Apply the tied linear transform.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor of shape (..., in_features), where in_features matches the second dimension of tied_module.weight.	required

Returns:

Type	Description
torch.Tensor	The output tensor of shape (..., out_features), where out_features matches the first dimension of tied_module.weight.