Skip to content

utorch.nn

utorch.nn

Neural network components, modules, layers, losses, and optimizers.

__all__ module-attribute 

__all__ = [
    "SGD",
    "CrossEntropyLoss",
    "Linear",
    "MSELoss",
    "Module",
    "Optimizer",
    "Parameter",
    "ReLU",
    "Sequential",
    "Sigmoid",
]

CrossEntropyLoss

Bases: Module

Cross-Entropy loss with logits using numerically stable LogSumExp.

Computes negative log-likelihood without requiring keepdims or ops changes

loss = - mean( sum( target * log_softmax(logits), axis=-1 ) )

forward 

forward(logits, target)

Compute negative log-likelihood without requiring keepdims or ops changes.

Source code in utorch/nn/loss.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def forward(self, logits: ut.Tensor, target: ut.Tensor) -> ut.Tensor:
    """Compute negative log-likelihood without requiring keepdims or ops changes."""
    # 1. Numerical stability: subtract max along class dimension
    # max_val has shape (64,)
    max_val = ut.Tensor(np.max(logits.value, axis=-1))
    # (64, 10).T - (64,) -> (10, 64).T -> (64, 10)
    shifted = (logits.transpose() - max_val).transpose()

    # 2. Sum along class axis (-1): produces shape (64,)
    sum_exp = shifted.exp().sum(axis=-1)
    log_sum_exp = sum_exp.log()  # Shape (64,)

    # 3. Transpose trick: (10, 64) - (64,) -> (10, 64).T -> (64, 10)
    log_probs = (shifted.transpose() - log_sum_exp).transpose()

    # 4. Negative log-likelihood against one-hot targets
    if target.shape == logits.shape:
        loss = (target * log_probs * (-1.0)).sum(axis=-1).mean()
    else:
        loss = log_probs.mean() * (-1.0)

    return loss

Linear 

Linear(in_features, out_features, bias=True, name=None)

Bases: Module

Fully connected linear transformation layer: Y = X @ W^T + b.

Weights are initialized using Kaiming (He) uniform initialization: W ~ Uniform(-1 / sqrt(in_features), 1 / sqrt(in_features))

Parameters:

Name Type Description Default
in_features int

Number of input features (D_in).

 required
out_features int

Number of output features (D_out).

 required
bias bool

If True, adds a learnable bias vector of shape (D_out,).

 True

Initialize a Linear layer.

Source code in utorch/nn/layers.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
def __init__(
    self,
    in_features: int,
    out_features: int,
    bias: bool = True,
    name: str | None = None,
) -> None:
    """Initialize a Linear layer."""
    super().__init__()
    self.in_features = in_features
    self.out_features = out_features

    self.__initialize(bias, in_features, name, out_features)

forward 

forward(x)

Forward linear transformation.

Source code in utorch/nn/layers.py 
56
57
58
59
60
61
def forward(self, x: ut.Tensor) -> ut.Tensor:
    """Forward linear transformation."""
    out = x @ self.weight.transpose()
    if self.bias is not None:
        out = out + self.bias
    return out

MSELoss

Bases: Module

Mean Squared Error (L2) loss: L = mean((pred - target)^2).

forward 

forward(pred, target)

Compute mean squared error (L2) loss.

Source code in utorch/nn/loss.py 
14
15
16
def forward(self, pred: ut.Tensor, target: ut.Tensor) -> ut.Tensor:
    """Compute mean squared error (L2) loss."""
    return ((pred - target) ** 2).mean()

Module

Base class for all neural network modules and layers.

Custom models should subclass Module and implement forward(). Modules automatically register child Module instances and Parameter attributes, enabling recursive traversal via parameters().

Examples:

import utorch as ut
import utorch.nn as nn

class TwoLayerMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

model = TwoLayerMLP()
print(len(model.parameters()))  # 4 (fc1.weight, fc1.bias, fc2.weight, fc2.bias)

__call__ 

__call__(*args, **kwargs)

Call self.forward(args, *kwargs), mimicking PyTorch's execution syntax.

Source code in utorch/nn/module.py 
43
44
45
def __call__(self, *args: Any, **kwargs: Any) -> Any:
    """Call self.forward(*args, **kwargs), mimicking PyTorch's execution syntax."""
    return self.forward(*args, **kwargs)

forward 

forward(*args, **kwargs)

Execute forward pass computation. Subclasses must override this.

Source code in utorch/nn/module.py 
39
40
41
def forward(self, *args: Any, **kwargs: Any) -> Any:
    """Execute forward pass computation. Subclasses must override this."""
    raise NotImplementedError("Subclasses of Module must implement forward().")

num_parameters 

num_parameters()

Return the total number of scalar parameters in the module.

Source code in utorch/nn/module.py 
86
87
88
def num_parameters(self) -> int:
    """Return the total number of scalar parameters in the module."""
    return sum(int(np.prod(p.shape)) for p in self.parameters())

parameters 

parameters()

Return an ordered list of all trainable Parameters in this module.

Recursively discovers parameters stored as direct attributes, inside lists/tuples, or within child Module submodules.

Source code in utorch/nn/module.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def parameters(self) -> list[Parameter]:
    """Return an ordered list of all trainable Parameters in this module.

    Recursively discovers parameters stored as direct attributes, inside
    lists/tuples, or within child `Module` submodules.
    """
    params: list[Parameter] = []

    for attr_name in dir(self):
        if attr_name.startswith("__"):
            continue
        val = getattr(self, attr_name)

        if isinstance(val, Parameter):
            params.append(val)
        elif isinstance(val, Module):
            params.extend(val.parameters())
        elif isinstance(val, list | tuple):
            for item in val:
                if isinstance(item, Parameter):
                    params.append(item)
                elif isinstance(item, Module):
                    params.extend(item.parameters())

    # Deduplicate preserving order
    seen: set[int] = set()
    unique_params: list[Parameter] = []
    for p in params:
        if id(p) not in seen:
            seen.add(id(p))
            unique_params.append(p)

    return unique_params

zero_grad 

zero_grad()

Reset gradients of all trainable parameters in this module to None.

Source code in utorch/nn/module.py 
81
82
83
84
def zero_grad(self) -> None:
    """Reset gradients of all trainable parameters in this module to None."""
    for p in self.parameters():
        p.zero_grad()

Optimizer 

Optimizer(params, lr)

Base class for all parameter optimizers.

Initialize a parameter optimizer.

Source code in utorch/nn/optim.py 
15
16
17
18
def __init__(self, params: Iterable[Parameter], lr: float) -> None:
    """Initialize a parameter optimizer."""
    self.params: list[Parameter] = list(params)
    self.lr: float = lr

step 

step()

Perform a single parameter update step.

Source code in utorch/nn/optim.py 
25
26
27
def step(self) -> None:
    """Perform a single parameter update step."""
    raise NotImplementedError

zero_grad 

zero_grad()

Reset gradients for all tracked parameters.

Source code in utorch/nn/optim.py 
20
21
22
23
def zero_grad(self) -> None:
    """Reset gradients for all tracked parameters."""
    for p in self.params:
        p.zero_grad()

Parameter 

Parameter(data, name=None)

Bases: Tensor

A Tensor designated as a trainable model parameter.

Parameters are automatically detected and registered by nn.Module instances when assigned as attributes, making them discoverable via module.parameters().

Parameters:

Name Type Description Default
data Any

Numerical data (NumPy array, nested list, or scalar).

 required
name str | None

Optional human-readable identifier for diagramming and debugging.

 None

Examples:

import numpy as np
import utorch as ut
from utorch.nn import Parameter

# Initialize a trainable weight matrix
w = Parameter(np.random.randn(3, 2), name="weight")
print(isinstance(w, ut.Tensor))  # True

Initialize a model parameter.

Source code in utorch/nn/parameter.py 
32
33
34
35
def __init__(self, data: Any, name: str | None = None) -> None:
    """Initialize a model parameter."""
    super().__init__(value=data, name=name)
    self.is_parameter: bool = True

ReLU

Bases: Module

Applies element-wise Rectified Linear Unit activation: max(0, x).

forward 

forward(x)

Forward ReLU transformation.

Source code in utorch/nn/layers.py 
67
68
69
def forward(self, x: ut.Tensor) -> ut.Tensor:
    """Forward ReLU transformation."""
    return x.relu()

SGD 

SGD(params, lr=0.01, momentum=0.0)

Bases: Optimizer

Stochastic Gradient Descent optimizer with optional classical momentum.

Parameters:

Name Type Description Default
params Iterable[Parameter]

Iterable of parameters to optimize.

 required
lr float

Learning rate.

 0.01
momentum float

Momentum factor (default: 0.0).

 0.0

Initialize SGD optimizer.

Source code in utorch/nn/optim.py 
39
40
41
42
43
44
45
46
47
48
def __init__(
    self,
    params: Iterable[Parameter],
    lr: float = 1e-2,
    momentum: float = 0.0,
) -> None:
    """Initialize SGD optimizer."""
    super().__init__(params, lr)
    self.momentum = momentum
    self.velocities = [np.zeros_like(p.value) for p in self.params]

step 

step()

Perform a single parameter update step.

Source code in utorch/nn/optim.py 
50
51
52
53
54
55
56
57
58
59
60
61
def step(self) -> None:
    """Perform a single parameter update step."""
    for i, p in enumerate(self.params):
        if p.grad is None:
            continue
        grad_val = p.grad.value

        if self.momentum > 0.0:
            self.velocities[i] = self.momentum * self.velocities[i] + grad_val
            p.value -= self.lr * self.velocities[i]
        else:
            p.value -= self.lr * grad_val

Sequential 

Sequential(*layers)

Bases: Module

Sequential container passing the output of each module into the next.

Parameters:

Name Type Description Default
*layers Module

Variable sequence of Module instances.

 ()

Examples:

import utorch.nn as nn

model = nn.Sequential(
    nn.Linear(2, 4),
    nn.ReLU(),
    nn.Linear(4, 1),
)

Initialize a Sequential container.

Source code in utorch/nn/layers.py 
 98
 99
100
101
def __init__(self, *layers: Module) -> None:
    """Initialize a Sequential container."""
    super().__init__()
    self.layers = list(layers)

__getitem__ 

__getitem__(index)

Return item at given index.

Source code in utorch/nn/layers.py 
109
110
111
def __getitem__(self, index: int) -> Module:
    """Return item at given index."""
    return self.layers[index]

__len__ 

__len__()

Return number of layers in sequential container.

Source code in utorch/nn/layers.py 
113
114
115
def __len__(self) -> int:
    """Return number of layers in sequential container."""
    return len(self.layers)

forward 

forward(x)

Forward pass through the sequential container.

Source code in utorch/nn/layers.py 
103
104
105
106
107
def forward(self, x: Any) -> Any:
    """Forward pass through the sequential container."""
    for layer in self.layers:
        x = layer(x)
    return x

Sigmoid

Bases: Module

Applies element-wise Sigmoid activation: 1 / (1 + exp(-x)).

forward 

forward(x)

Forward Sigmoid transformation.

Source code in utorch/nn/layers.py 
75
76
77
def forward(self, x: ut.Tensor) -> ut.Tensor:
    """Forward Sigmoid transformation."""
    return x.sigmoid()

layers

Neural network layers, activation modules, and sequential containers.

Linear 

Linear(in_features, out_features, bias=True, name=None)

Bases: Module

Fully connected linear transformation layer: Y = X @ W^T + b.

Weights are initialized using Kaiming (He) uniform initialization: W ~ Uniform(-1 / sqrt(in_features), 1 / sqrt(in_features))

Parameters:

Name Type Description Default
in_features int

Number of input features (D_in).

 required
out_features int

Number of output features (D_out).

 required
bias bool

If True, adds a learnable bias vector of shape (D_out,).

 True

Initialize a Linear layer.

Source code in utorch/nn/layers.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
def __init__(
    self,
    in_features: int,
    out_features: int,
    bias: bool = True,
    name: str | None = None,
) -> None:
    """Initialize a Linear layer."""
    super().__init__()
    self.in_features = in_features
    self.out_features = out_features

    self.__initialize(bias, in_features, name, out_features)
forward 
forward(x)

Forward linear transformation.

Source code in utorch/nn/layers.py 
56
57
58
59
60
61
def forward(self, x: ut.Tensor) -> ut.Tensor:
    """Forward linear transformation."""
    out = x @ self.weight.transpose()
    if self.bias is not None:
        out = out + self.bias
    return out

ReLU

Bases: Module

Applies element-wise Rectified Linear Unit activation: max(0, x).

forward 
forward(x)

Forward ReLU transformation.

Source code in utorch/nn/layers.py 
67
68
69
def forward(self, x: ut.Tensor) -> ut.Tensor:
    """Forward ReLU transformation."""
    return x.relu()

Sequential 

Sequential(*layers)

Bases: Module

Sequential container passing the output of each module into the next.

Parameters:

Name Type Description Default
*layers Module

Variable sequence of Module instances.

 ()

Examples:

import utorch.nn as nn

model = nn.Sequential(
    nn.Linear(2, 4),
    nn.ReLU(),
    nn.Linear(4, 1),
)

Initialize a Sequential container.

Source code in utorch/nn/layers.py 
 98
 99
100
101
def __init__(self, *layers: Module) -> None:
    """Initialize a Sequential container."""
    super().__init__()
    self.layers = list(layers)
__getitem__ 
__getitem__(index)

Return item at given index.

Source code in utorch/nn/layers.py 
109
110
111
def __getitem__(self, index: int) -> Module:
    """Return item at given index."""
    return self.layers[index]
__len__ 
__len__()

Return number of layers in sequential container.

Source code in utorch/nn/layers.py 
113
114
115
def __len__(self) -> int:
    """Return number of layers in sequential container."""
    return len(self.layers)
forward 
forward(x)

Forward pass through the sequential container.

Source code in utorch/nn/layers.py 
103
104
105
106
107
def forward(self, x: Any) -> Any:
    """Forward pass through the sequential container."""
    for layer in self.layers:
        x = layer(x)
    return x

Sigmoid

Bases: Module

Applies element-wise Sigmoid activation: 1 / (1 + exp(-x)).

forward 
forward(x)

Forward Sigmoid transformation.

Source code in utorch/nn/layers.py 
75
76
77
def forward(self, x: ut.Tensor) -> ut.Tensor:
    """Forward Sigmoid transformation."""
    return x.sigmoid()

loss

Loss functions for training neural networks.

CrossEntropyLoss

Bases: Module

Cross-Entropy loss with logits using numerically stable LogSumExp.

Computes negative log-likelihood without requiring keepdims or ops changes

loss = - mean( sum( target * log_softmax(logits), axis=-1 ) )

forward 
forward(logits, target)

Compute negative log-likelihood without requiring keepdims or ops changes.

Source code in utorch/nn/loss.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def forward(self, logits: ut.Tensor, target: ut.Tensor) -> ut.Tensor:
    """Compute negative log-likelihood without requiring keepdims or ops changes."""
    # 1. Numerical stability: subtract max along class dimension
    # max_val has shape (64,)
    max_val = ut.Tensor(np.max(logits.value, axis=-1))
    # (64, 10).T - (64,) -> (10, 64).T -> (64, 10)
    shifted = (logits.transpose() - max_val).transpose()

    # 2. Sum along class axis (-1): produces shape (64,)
    sum_exp = shifted.exp().sum(axis=-1)
    log_sum_exp = sum_exp.log()  # Shape (64,)

    # 3. Transpose trick: (10, 64) - (64,) -> (10, 64).T -> (64, 10)
    log_probs = (shifted.transpose() - log_sum_exp).transpose()

    # 4. Negative log-likelihood against one-hot targets
    if target.shape == logits.shape:
        loss = (target * log_probs * (-1.0)).sum(axis=-1).mean()
    else:
        loss = log_probs.mean() * (-1.0)

    return loss

MSELoss

Bases: Module

Mean Squared Error (L2) loss: L = mean((pred - target)^2).

forward 
forward(pred, target)

Compute mean squared error (L2) loss.

Source code in utorch/nn/loss.py 
14
15
16
def forward(self, pred: ut.Tensor, target: ut.Tensor) -> ut.Tensor:
    """Compute mean squared error (L2) loss."""
    return ((pred - target) ** 2).mean()

module

Base Module class for neural network layers and architectures.

Module

Base class for all neural network modules and layers.

Custom models should subclass Module and implement forward(). Modules automatically register child Module instances and Parameter attributes, enabling recursive traversal via parameters().

Examples:

import utorch as ut
import utorch.nn as nn

class TwoLayerMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

model = TwoLayerMLP()
print(len(model.parameters()))  # 4 (fc1.weight, fc1.bias, fc2.weight, fc2.bias)
__call__ 
__call__(*args, **kwargs)

Call self.forward(args, *kwargs), mimicking PyTorch's execution syntax.

Source code in utorch/nn/module.py 
43
44
45
def __call__(self, *args: Any, **kwargs: Any) -> Any:
    """Call self.forward(*args, **kwargs), mimicking PyTorch's execution syntax."""
    return self.forward(*args, **kwargs)
forward 
forward(*args, **kwargs)

Execute forward pass computation. Subclasses must override this.

Source code in utorch/nn/module.py 
39
40
41
def forward(self, *args: Any, **kwargs: Any) -> Any:
    """Execute forward pass computation. Subclasses must override this."""
    raise NotImplementedError("Subclasses of Module must implement forward().")
num_parameters 
num_parameters()

Return the total number of scalar parameters in the module.

Source code in utorch/nn/module.py 
86
87
88
def num_parameters(self) -> int:
    """Return the total number of scalar parameters in the module."""
    return sum(int(np.prod(p.shape)) for p in self.parameters())
parameters 
parameters()

Return an ordered list of all trainable Parameters in this module.

Recursively discovers parameters stored as direct attributes, inside lists/tuples, or within child Module submodules.

Source code in utorch/nn/module.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def parameters(self) -> list[Parameter]:
    """Return an ordered list of all trainable Parameters in this module.

    Recursively discovers parameters stored as direct attributes, inside
    lists/tuples, or within child `Module` submodules.
    """
    params: list[Parameter] = []

    for attr_name in dir(self):
        if attr_name.startswith("__"):
            continue
        val = getattr(self, attr_name)

        if isinstance(val, Parameter):
            params.append(val)
        elif isinstance(val, Module):
            params.extend(val.parameters())
        elif isinstance(val, list | tuple):
            for item in val:
                if isinstance(item, Parameter):
                    params.append(item)
                elif isinstance(item, Module):
                    params.extend(item.parameters())

    # Deduplicate preserving order
    seen: set[int] = set()
    unique_params: list[Parameter] = []
    for p in params:
        if id(p) not in seen:
            seen.add(id(p))
            unique_params.append(p)

    return unique_params
zero_grad 
zero_grad()

Reset gradients of all trainable parameters in this module to None.

Source code in utorch/nn/module.py 
81
82
83
84
def zero_grad(self) -> None:
    """Reset gradients of all trainable parameters in this module to None."""
    for p in self.parameters():
        p.zero_grad()

optim

Optimization algorithms for updating model parameters.

Optimizer 

Optimizer(params, lr)

Base class for all parameter optimizers.

Initialize a parameter optimizer.

Source code in utorch/nn/optim.py 
15
16
17
18
def __init__(self, params: Iterable[Parameter], lr: float) -> None:
    """Initialize a parameter optimizer."""
    self.params: list[Parameter] = list(params)
    self.lr: float = lr
step 
step()

Perform a single parameter update step.

Source code in utorch/nn/optim.py 
25
26
27
def step(self) -> None:
    """Perform a single parameter update step."""
    raise NotImplementedError
zero_grad 
zero_grad()

Reset gradients for all tracked parameters.

Source code in utorch/nn/optim.py 
20
21
22
23
def zero_grad(self) -> None:
    """Reset gradients for all tracked parameters."""
    for p in self.params:
        p.zero_grad()

SGD 

SGD(params, lr=0.01, momentum=0.0)

Bases: Optimizer

Stochastic Gradient Descent optimizer with optional classical momentum.

Parameters:

Name Type Description Default
params Iterable[Parameter]

Iterable of parameters to optimize.

 required
lr float

Learning rate.

 0.01
momentum float

Momentum factor (default: 0.0).

 0.0

Initialize SGD optimizer.

Source code in utorch/nn/optim.py 
39
40
41
42
43
44
45
46
47
48
def __init__(
    self,
    params: Iterable[Parameter],
    lr: float = 1e-2,
    momentum: float = 0.0,
) -> None:
    """Initialize SGD optimizer."""
    super().__init__(params, lr)
    self.momentum = momentum
    self.velocities = [np.zeros_like(p.value) for p in self.params]
step 
step()

Perform a single parameter update step.

Source code in utorch/nn/optim.py 
50
51
52
53
54
55
56
57
58
59
60
61
def step(self) -> None:
    """Perform a single parameter update step."""
    for i, p in enumerate(self.params):
        if p.grad is None:
            continue
        grad_val = p.grad.value

        if self.momentum > 0.0:
            self.velocities[i] = self.momentum * self.velocities[i] + grad_val
            p.value -= self.lr * self.velocities[i]
        else:
            p.value -= self.lr * grad_val

parameter

Trainable network parameter abstraction.

Parameter 

Parameter(data, name=None)

Bases: Tensor

A Tensor designated as a trainable model parameter.

Parameters are automatically detected and registered by nn.Module instances when assigned as attributes, making them discoverable via module.parameters().

Parameters:

Name Type Description Default
data Any

Numerical data (NumPy array, nested list, or scalar).

 required
name str | None

Optional human-readable identifier for diagramming and debugging.

 None

Examples:

import numpy as np
import utorch as ut
from utorch.nn import Parameter

# Initialize a trainable weight matrix
w = Parameter(np.random.randn(3, 2), name="weight")
print(isinstance(w, ut.Tensor))  # True

Initialize a model parameter.

Source code in utorch/nn/parameter.py 
32
33
34
35
def __init__(self, data: Any, name: str | None = None) -> None:
    """Initialize a model parameter."""
    super().__init__(value=data, name=name)
    self.is_parameter: bool = True