Skip to main content

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/MilesONerd/neurenix/llms.txt

Use this file to discover all available pages before exploring further.

​
Overview

Neurenix provides a comprehensive set of neural network layers for building deep learning models.

​
Linear Layer

class Linear(Module):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        dtype: Optional[DType] = None,
        device: Optional[Device] = None,
    )
Linear transformation: y = xW^T + b

​
Parameters

​
in_features
int
required
Size of each input sample.
​
out_features
int
required
Size of each output sample.
​
bias
bool
default:"True"
If True, adds a learnable bias to the output.
​
dtype
Optional[DType]
Data type of the parameters.
​
device
Optional[Device]
Device to store the parameters on.

​
Attributes

  • weight: Learnable weights of shape (out_features, in_features)
  • bias: Learnable bias of shape (out_features,) (if bias=True)

​
Example

import neurenix as nx

# Create a linear layer
layer = nx.Linear(784, 128)

# Forward pass
x = nx.Tensor.randn((32, 784))
output = layer(x)
print(output.shape)  # (32, 128)

# Without bias
layer_no_bias = nx.Linear(128, 10, bias=False)

​
Convolutional Layers

​
Conv1d

class Conv1d(Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        stride: int = 1,
        padding: int = 0,
        dilation: int = 1,
        groups: int = 1,
        bias: bool = True,
    )
1D convolution layer.
​
in_channels
int
required
Number of input channels.
​
out_channels
int
required
Number of output channels.
​
kernel_size
int
required
Size of the convolving kernel.
​
stride
int
default:"1"
Stride of the convolution.
​
padding
int
default:"0"
Zero-padding added to both sides of the input.

​
Conv2d

class Conv2d(Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Union[int, Tuple[int, int]],
        stride: Union[int, Tuple[int, int]] = 1,
        padding: Union[int, Tuple[int, int]] = 0,
        dilation: Union[int, Tuple[int, int]] = 1,
        groups: int = 1,
        bias: bool = True,
    )
2D convolution layer for image processing.

​
Conv3d

class Conv3d(Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Union[int, Tuple[int, int, int]],
        stride: Union[int, Tuple[int, int, int]] = 1,
        padding: Union[int, Tuple[int, int, int]] = 0,
        dilation: Union[int, Tuple[int, int, int]] = 1,
        groups: int = 1,
        bias: bool = True,
    )
3D convolution layer for video or volumetric data.

​
Example

# 2D convolution for images
conv = nx.Conv2d(
    in_channels=3,
    out_channels=64,
    kernel_size=3,
    stride=1,
    padding=1
)

# Input: batch of RGB images
x = nx.Tensor.randn((32, 3, 224, 224))
output = conv(x)
print(output.shape)  # (32, 64, 224, 224)

# 1D convolution for sequences
conv1d = nx.Conv1d(
    in_channels=100,
    out_channels=256,
    kernel_size=5,
    stride=1
)

seq = nx.Tensor.randn((16, 100, 50))  # (batch, channels, length)
output = conv1d(seq)
print(output.shape)  # (16, 256, 46)

​
Recurrent Layers

​
RNN

class RNN(Module):
    def __init__(
        self,
        input_size: int,
        hidden_size: int,
        num_layers: int = 1,
        nonlinearity: str = 'tanh',
        bias: bool = True,
        batch_first: bool = False,
        dropout: float = 0.0,
        bidirectional: bool = False,
    )
Vanilla recurrent neural network.

​
LSTM

class LSTM(Module):
    def __init__(
        self,
        input_size: int,
        hidden_size: int,
        num_layers: int = 1,
        bias: bool = True,
        batch_first: bool = False,
        dropout: float = 0.0,
        bidirectional: bool = False,
    )
Long Short-Term Memory network.
​
input_size
int
required
The number of expected features in the input.
​
hidden_size
int
required
The number of features in the hidden state.
​
num_layers
int
default:"1"
Number of recurrent layers.
​
batch_first
bool
default:"False"
If True, input and output tensors are provided as (batch, seq, feature).
​
bidirectional
bool
default:"False"
If True, becomes a bidirectional LSTM.

​
GRU

class GRU(Module):
    def __init__(
        self,
        input_size: int,
        hidden_size: int,
        num_layers: int = 1,
        bias: bool = True,
        batch_first: bool = False,
        dropout: float = 0.0,
        bidirectional: bool = False,
    )
Gated Recurrent Unit.

​
Example

# LSTM for sequence modeling
lstm = nx.LSTM(
    input_size=128,
    hidden_size=256,
    num_layers=2,
    batch_first=True,
    bidirectional=True
)

# Input: (batch, sequence_length, features)
x = nx.Tensor.randn((32, 50, 128))
output, (hidden, cell) = lstm(x)

# output shape: (32, 50, 512) - 512 because bidirectional
print(output.shape)

# GRU for language modeling
gru = nx.GRU(
    input_size=300,  # word embedding dimension
    hidden_size=512,
    num_layers=3,
    batch_first=True,
    dropout=0.2
)

embeddings = nx.Tensor.randn((16, 100, 300))
output, hidden = gru(embeddings)

​
Pooling Layers

​
MaxPool2d

class MaxPool2d(Module):
    def __init__(
        self,
        kernel_size: Union[int, Tuple[int, int]],
        stride: Optional[Union[int, Tuple[int, int]]] = None,
        padding: Union[int, Tuple[int, int]] = 0,
        dilation: Union[int, Tuple[int, int]] = 1,
    )
Max pooling over a 2D input.
​
kernel_size
Union[int, Tuple[int, int]]
required
Size of the pooling window.
​
stride
Optional[Union[int, Tuple[int, int]]]
Stride of the pooling window. Default is kernel_size.

​
Example

# Max pooling
maxpool = nx.MaxPool2d(kernel_size=2, stride=2)

x = nx.Tensor.randn((32, 64, 224, 224))
output = maxpool(x)
print(output.shape)  # (32, 64, 112, 112)

​
Regularization Layers

​
Dropout

class Dropout(Module):
    def __init__(self, p: float = 0.5, inplace: bool = False)
Randomly zeroes elements during training.
​
p
float
default:"0.5"
Probability of an element to be zeroed.
​
inplace
bool
default:"False"
If True, will do this operation in-place.

​
Example

dropout = nx.Dropout(p=0.5)

model.train()  # Dropout active
x = nx.Tensor.randn((32, 128))
output = dropout(x)  # Some elements zeroed

model.eval()  # Dropout inactive
output = dropout(x)  # No elements zeroed

​
Container Layers

​
Sequential

class Sequential(Module):
    def __init__(self, *args)
Sequential container for chaining layers.

​
Example

# Build a simple CNN
model = nx.Sequential(
    nx.Conv2d(3, 64, kernel_size=3, padding=1),
    nx.ReLU(),
    nx.MaxPool2d(2),
    nx.Conv2d(64, 128, kernel_size=3, padding=1),
    nx.ReLU(),
    nx.MaxPool2d(2),
    nx.Linear(128 * 56 * 56, 10)
)

# Forward pass
x = nx.Tensor.randn((1, 3, 224, 224))
output = model(x)
print(output.shape)  # (1, 10)

​
Complete Example

import neurenix as nx

class ResNetBlock(nx.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv1 = nx.Conv2d(channels, channels, 3, padding=1)
        self.conv2 = nx.Conv2d(channels, channels, 3, padding=1)
        self.relu = nx.ReLU()
    
    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.relu(out)
        out = self.conv2(out)
        out = out + residual
        out = self.relu(out)
        return out

class MyNetwork(nx.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nx.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
        self.pool = nx.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.res_blocks = nx.Sequential(
            ResNetBlock(64),
            ResNetBlock(64),
            ResNetBlock(64)
        )
        self.fc = nx.Linear(64 * 56 * 56, 1000)
        self.dropout = nx.Dropout(0.5)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.pool(x)
        x = self.res_blocks(x)
        x = x.reshape(x.shape[0], -1)
        x = self.dropout(x)
        x = self.fc(x)
        return x

model = MyNetwork()
print(nx.utils.model_summary(model))

Build docs developers (and LLMs) love

Get started for freeTalk to us