### Install x-mlps-pytorch

Source: https://github.com/lucidrains/x-mlps-pytorch/blob/main/README.md

Install the package using pip.

```bash
$ pip install x-mlps-pytorch
```

--------------------------------

### Initialize and use Ensemble

Source: https://context7.com/lucidrains/x-mlps-pytorch/llms.txt

Wrap any PyTorch module to create an ensemble of models for efficient parallel forward passes using vmap. Parameters include the base network, ensemble size, and initialization standard deviation. Supports forwarding all members, specific members, single members, and averaging.

```python
import torch
from x_mlps_pytorch import MLP, Ensemble

# Create base network
base_net = MLP(64, 128, 32)

# Create ensemble of 5 networks
ensemble = Ensemble(
    net=base_net,
    ensemble_size=5,
    init_std_dev=0.02
)

# Forward all ensemble members
x = torch.randn(16, 64)
outputs = ensemble(x)  # (5, 16, 32)

# Forward specific ensemble members
outputs_subset = ensemble(x, ids=[0, 2, 4])  # (3, 16, 32)

# Forward single member
output_single = ensemble.forward_one(x, id=0)  # (16, 32)

# Get averaged model from ensemble members
averaged_model = ensemble.get_one(
    indices=[0, 1, 2],
    weights=[0.5, 0.3, 0.2]
)
output_avg = averaged_model(x)  # (16, 32)

# Per-batch-sample forward (different ensemble member per batch item)
# Requires batch_size == ensemble_size
x_per_sample = torch.randn(5, 64)
outputs_per_sample = ensemble(x_per_sample, each_batch_sample=True)
```

--------------------------------

### Initialize and Use MLPs

Source: https://github.com/lucidrains/x-mlps-pytorch/blob/main/README.md

Instantiate MLP models and perform forward passes with input tensors.

```python
import torch
from x_mlps_pytorch import MLP

actor = MLP(10, 16, 5)

critic = MLP(10, 32, 16, 1)

state = torch.randn(10)

action_logits = actor(state) # (5,)

values = critic(state) # (1,)
```

--------------------------------

### Initialize and use GroupedFeedforwards

Source: https://context7.com/lucidrains/x-mlps-pytorch/llms.txt

Create parallel transformer-style feedforward blocks using grouped convolutions. Parameters include dimension, depth, number of groups, expansion factor, and optional final normalization. Can also include input and output projections.

```python
import torch
from x_mlps_pytorch import GroupedFeedforwards

# 4 parallel feedforward stacks
grouped_ff = GroupedFeedforwards(
    dim=256,
    depth=4,
    groups=4,
    expansion_factor=4.,
    final_norm=True
)

# With input/output projections
grouped_ff_proj = GroupedFeedforwards(
    dim=256,
    depth=3,
    dim_in=64,
    dim_out=32,
    groups=8,
    activation=torch.nn.GELU(),
    squeeze_if_one_group=False
)

# Forward pass
x = torch.randn(16, 64)
output = grouped_ff_proj(x)  # (16, 8, 32)
```

--------------------------------

### Initialize and use nGroupedFeedforwards

Source: https://context7.com/lucidrains/x-mlps-pytorch/llms.txt

Implement parallel normalized feedforward networks using hyperspherical normalization and grouped convolutions. Key parameters include dimension, depth, number of groups, input/output dimensions, and options for preserving input magnitude and squeezing single groups.

```python
import torch
from x_mlps_pytorch import nGroupedFeedforwards

# 4 parallel hyperspherical feedforward stacks
ngff = nGroupedFeedforwards(
    dim=512,
    depth=4,
    groups=4,
    dim_in=128,
    dim_out=128,
    input_preserve_magnitude=True,
    squeeze_if_one_group=False
)

# Forward pass
x = torch.randn(2, 128)
output = ngff(x)  # (2, 4, 128)

# With single group and squeeze
ngff_single = nGroupedFeedforwards(
    dim=512,
    depth=4,
    groups=1,
    dim_in=128,
    dim_out=128,
    input_preserve_magnitude=True,
    squeeze_if_one_group=True
)
output_squeezed = ngff_single(x)  # (2, 128)
```

--------------------------------

### Initialize and use ResidualNormedMLP

Source: https://context7.com/lucidrains/x-mlps-pytorch/llms.txt

Create deep MLPs with residual connections every N layers, suitable for goal-conditioned reinforcement learning. Parameters include dimension, depth, residual frequency, input/output dimensions, activation function, and normalization options. Supports stabilization and skip connections to output.

```python
import torch
from x_mlps_pytorch import ResidualNormedMLP

# Deep residual MLP: 32 layers with residual every 4
deep_mlp = ResidualNormedMLP(
    dim=256,
    depth=32,
    residual_every=4,
    dim_in=64,
    dim_out=32,
    activation=torch.nn.SiLU(),  # SiLU is important per paper ablation
    use_rmsnorm=True,
    final_norm=True
)

# With Keel post-LayerNorm stabilization (Chen et al.)
stable_mlp = ResidualNormedMLP(
    dim=256,
    depth=16,
    residual_every=4,
    keel_post_ln=True,
    keel_residual_scale=4  # defaults to num_blocks
)

# Auto-compression network (Dorovatas et al.) - skip connections to output
auto_compress_mlp = ResidualNormedMLP(
    dim=256,
    depth=16,
    residual_every=4,
    skip_to_output=True  # sum all block outputs
)

# Forward pass with concatenated inputs
state = torch.randn(32, 32)
goal = torch.randn(32, 32)
output = deep_mlp([state, goal])  # (32, 32)
```

--------------------------------

### Apply Noisable to Model Parameters

Source: https://context7.com/lucidrains/x-mlps-pytorch/llms.txt

Wraps a model to inject noise into parameters for exploration or uncertainty estimation. Supports seed-based reproducible noise, tensor-based noise, and in-place modifications.

```python
import torch
from x_mlps_pytorch import MLP, Noisable, with_seed

# Create base model and wrap with Noisable
base_model = MLP(64, 128, 32)
noisable = Noisable(
    model=base_model,
    noise_scale=0.01,
    low_rank=8  # optional low-rank noise
)

# Get parameter names for noise specification
param_names = list(base_model.named_parameters())
# Example: [('layers.0.0.weight', ...), ('layers.0.0.bias', ...), ...]

# Forward with seed-based noise (reproducible)
x = torch.randn(16, 64)
noise_spec = {
    'layers.0.0.weight': 42,  # seed
    'layers.1.weight': 123,
}
output_noised = noisable(x, noise_for_params=noise_spec)

# Forward with tensor noise
noise_tensor = torch.randn(128, 64) * 0.01
noise_spec_tensor = {
    'layers.0.0.weight': noise_tensor,
}
output_noised = noisable(x, noise_for_params=noise_spec_tensor)

# Seed with custom scale: (seed, scale)
noise_spec_scaled = {
    'layers.0.0.weight': (42, 0.05),  # seed=42, scale=0.05
}
output_scaled = noisable(x, noise_for_params=noise_spec_scaled)

# Temporarily add noise in-place (context manager)
with noisable.temp_add_noise_(noise_for_params={'layers.0.0.weight': 42}):
    output_temp = noisable.model(x)  # model has noise
# noise automatically removed

# Permanently add noise in-place
noisable.add_noise_(noise_for_params={'layers.0.0.weight': 42})
```

--------------------------------

### Initialize and use GroupedMLP

Source: https://context7.com/lucidrains/x-mlps-pytorch/llms.txt

Instantiate GroupedMLP for parallel MLPs using grouped convolutions. Supports multiple groups or a single group with optional squeezing. Input dimension, hidden dimension, output dimension, and number of groups are key parameters.

```python
import torch
from x_mlps_pytorch import GroupedMLP

# 4 parallel MLPs: each 64 -> 128 -> 32
grouped_mlp = GroupedMLP(
    64, 128, 32,
    groups=4,
    activation=torch.nn.ReLU()
)

# Input: batch=16, dim=64
x = torch.randn(16, 64)

# Output: batch=16, groups=4, dim=32
output = grouped_mlp(x)
print(output.shape)  # torch.Size([16, 4, 32])

# With squeeze for single group (behaves like regular MLP)
single_group = GroupedMLP(
    64, 128, 32,
    groups=1,
    squeeze_if_one_group=True
)
output_squeezed = single_group(x)  # (16, 32) instead of (16, 1, 32)
```

--------------------------------

### Basic MLP Implementation in PyTorch

Source: https://context7.com/lucidrains/x-mlps-pytorch/llms.txt

Use the MLP class for fundamental multi-layer perceptrons. It supports configurable dimensions, activations, and bias, and can handle both single and concatenated multi-tensor inputs.

```python
import torch
from x_mlps_pytorch import MLP

# Basic MLP: 10 input -> 16 hidden -> 5 output
actor = MLP(10, 16, 5)

# Deeper MLP: 10 -> 32 -> 16 -> 1
critic = MLP(10, 32, 16, 1)

# With custom activation and no bias
custom_mlp = MLP(64, 128, 64, 32, activation=torch.nn.SiLU(), bias=False)

# Activate the last layer (useful for bounded outputs)
bounded_mlp = MLP(10, 32, 10, activation=torch.nn.Tanh(), activate_last=True)

# Forward pass
state = torch.randn(32, 10)  # batch of 32, dim 10
action_logits = actor(state)  # (32, 5)
values = critic(state)  # (32, 1)

# MLP also accepts list/tuple inputs (automatically concatenates)
state1 = torch.randn(32, 5)
state2 = torch.randn(32, 5)
combined_output = actor([state1, state2])  # (32, 5)
```

--------------------------------

### Configure LatentConditionedFeedforwards

Source: https://context7.com/lucidrains/x-mlps-pytorch/llms.txt

Implements feedforward blocks conditioned on latent variables using adaptive normalization. Supports both standard and sequential input data.

```python
import torch
from x_mlps_pytorch.ff_with_latent import LatentConditionedFeedforwards

# Feedforward conditioned on latent (e.g., skill embedding)
conditioned_ff = LatentConditionedFeedforwards(
    dim=256,
    depth=4,
    dim_in=64,
    dim_out=32,
    dim_latent=16,  # latent/conditioning dimension
    latent_mlp=True,  # process latent through MLP first
    expansion_factor=4.,
    final_norm=True
)

# Forward pass
state = torch.randn(32, 64)  # observations
latent = torch.randn(32, 16)  # skill/gene embedding
output = conditioned_ff(state, latent)  # (32, 32)

# Works with sequential data too
state_seq = torch.randn(32, 10, 64)  # batch, seq_len, dim
output_seq = conditioned_ff(state_seq, latent)  # (32, 10, 32)
```

--------------------------------

### Utilize Custom Activation Functions

Source: https://context7.com/lucidrains/x-mlps-pytorch/llms.txt

Provides specialized activation functions like ReLU squared variants and Sugar activations. These can be integrated directly into MLP architectures.

```python
import torch
from x_mlps_pytorch.activations import ReluSquared, BSiLU, NeLU, Sugar, ReluNelu

# ReLU squared: ReLU(x)^2
relu_sq = ReluSquared(signed=False)
x = torch.randn(32, 64)
out = relu_sq(x)  # all non-negative

# Signed ReLU squared: sign(x) * ReLU(x)^2
relu_sq_signed = ReluSquared(signed=True)
out_signed = relu_sq_signed(x)  # preserves sign

# BSiLU (Bounded SiLU) from "The Resurrection of ReLU"
bsilu = BSiLU(alpha=1.67)
out_bsilu = bsilu(x)

# NeLU (Negative only activation)
nelu = NeLU(alpha=0.05)
out_nelu = nelu(x)

# Sugar activation with straight-through for negative region
# ReluNelu: ReLU forward, NeLU gradients for negative inputs
relu_nelu = ReluNelu(alpha=0.05)
out_relu_nelu = relu_nelu(x)

# Use custom activation in MLP
from x_mlps_pytorch import MLP
mlp = MLP(64, 128, 32, activation=ReluNelu())
```

--------------------------------

### Implement GradientDropout

Source: https://context7.com/lucidrains/x-mlps-pytorch/llms.txt

Applies dropout exclusively to gradients during backpropagation, leaving the forward pass unaffected. Useful for regularization without altering inference behavior.

```python
import torch
from x_mlps_pytorch import GradientDropout

# 50% gradient dropout
grad_dropout = GradientDropout(prob=0.5)

# Forward pass (unchanged)
t = torch.randn(50, requires_grad=True)
out = grad_dropout(t)

# Values are identical
assert torch.allclose(t, out)

# Backward pass has ~50% gradients dropped
out.sum().backward()
print(t.grad)  # approximately half of gradients are zero

# Use in a model
class RegularizedMLP(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(64, 128)
        self.grad_dropout = GradientDropout(0.3)
        self.fc2 = torch.nn.Linear(128, 32)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.grad_dropout(x)  # only affects gradients
        return self.fc2(x)
```

--------------------------------

### Citations

Source: https://github.com/lucidrains/x-mlps-pytorch/blob/main/README.md

BibTeX citations for the research papers referenced in the project.

```bibtex
@article{So2021PrimerSF,
    title   = {Primer: Searching for Efficient Transformers for Language Modeling},
    author  = {David R. So and Wojciech Ma'nke and Hanxiao Liu and Zihang Dai and Noam M. Shazeer and Quoc V. Le},
    journal = {ArXiv},
    year    = {2021},
    volume  = {abs/2109.08668},
    url     = {https://api.semanticscholar.org/CorpusID:237563187}
}
```

```bibtex
@article{Zhang2024ReLU2WD,
    title   = {ReLU2 Wins: Discovering Efficient Activation Functions for Sparse LLMs},
    author  = {Zhengyan Zhang and Yixin Song and Guanghui Yu and Xu Han and Yankai Lin and Chaojun Xiao and Chenyang Song and Zhiyuan Liu and Zeyu Mi and Maosong Sun},
    journal = {ArXiv},
    year    = {2024},
    volume  = {abs/2402.03804},
    url     = {https://api.semanticscholar.org/CorpusID:267499856}
}
```

```bibtex
@inproceedings{Horuz2025TheRO,
    title   = {The Resurrection of the ReLU},
    author  = {Cocsku Can Horuz and Geoffrey Kasenbacher and Saya Higuchi and Sebastian Kairat and Jendrik Stoltz and Moritz Pesl and Bernhard A. Moser and Christoph Linse and Thomas Martinetz and Sebastian Otte},
    year    = {2025},
    url     = {https://api.semanticscholar.org/CorpusID:278959515}
}
```

```bibtex
@article{Loshchilov2024nGPTNT,
    title   = {nGPT: Normalized Transformer with Representation Learning on the Hypersphere},
    author  = {Ilya Loshchilov and Cheng-Ping Hsieh and Simeng Sun and Boris Ginsburg},
    journal = {ArXiv},
    year    = {2024},
    volume  = {abs/2410.01131},
    url     = {https://api.semanticscholar.org/CorpusID:273026160}
}
```

```bibtex
@article{Lee2025HypersphericalNF,
    title   = {Hyperspherical Normalization for Scalable Deep Reinforcement Learning},
    author  = {Hojoon Lee and Youngdo Lee and Takuma Seno and Donghu Kim and Peter Stone and Jaegul Choo},
    journal = {ArXiv},
    year    = {2025},
    volume  = {abs/2502.15280},
    url     = {https://api.semanticscholar.org/CorpusID:276558261}
}
```

```bibtex
@inproceedings{wang2025,
    title   = {1000 Layer Networks for Self-Supervised {RL}: Scaling Depth Can Enable New Goal-Reaching Capabilities},
    author  = {Kevin Wang and Ishaan Javali and Micha{\l} Bortkiewicz and Tomasz Trzcinski and Benjamin Eysenbach},
    booktitle = {The Thirty-ninth Annual Conference on Neural Information Processing Systems},
    year    = {2025},
    url     = {https://openreview.net/forum?id=s0JVsx3bx1}
}
```

```bibtex
@misc{chen2026postlayernormbackstableexpressive,
    title   = {Post-LayerNorm Is Back: Stable, ExpressivE, and Deep}, 
    author  = {Chen Chen and Lai Wei},
    year    = {2026},
    eprint  = {2601.19895},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url     = {https://arxiv.org/abs/2601.19895}, 
}
```

```bibtex
@inproceedings{dorovatas2025autocompressing,
    title  = {Auto-Compressing Networks},
    author = {Vaggelis Dorovatas and Georgios Paraskevopoulos and Alexandros Potamianos},
    booktitle = {The Thirty-ninth Annual Conference on Neural Information Processing Systems},
    year    = {2025},
    url     = {https://openreview.net/forum?id=eIDa6pd9iQ}
}
```

```bibtex
@inproceedings{Lin2025ContinualLV,
    title    = {Continual Learning via Sparse Memory Finetuning},
    author   = {Jessy Lin and Luke S. Zettlemoyer and Gargi Ghosh and Wen-tau Yih and Aram H. Markosyan and Vincent-Pierre Berges and Barlas Ouguz},
    year     = {2025},
    url      = {https://api.semanticscholar.org/CorpusID:282203348},
    blog_url = {https://jessylin.com/2025/10/20/continual-learning/}
}
```

--------------------------------

### Factory Function for Creating MLPs

Source: https://context7.com/lucidrains/x-mlps-pytorch/llms.txt

Utilize create_mlp for convenient MLP generation with uniform hidden layer sizes. Specify depth, hidden dimension, and optional input/output dimensions.

```python
import torch
from x_mlps_pytorch import MLP, create_mlp

# Create MLP with depth=3 (4 layers total) and hidden dim=256
# Results in: 256 -> 256 -> 256 -> 256
mlp = create_mlp(dim=256, depth=3)

# With different input and output dimensions
# Results in: 64 -> 128 -> 128 -> 128 -> 32
mlp_with_io = create_mlp(
    dim=128,
    depth=2,
    dim_in=64,
    dim_out=32,
    activation=torch.nn.GELU()
)

# Forward pass
x = torch.randn(16, 64)
output = mlp_with_io(x)  # (16, 32)
```

--------------------------------

### Transformer-Style Feedforward Blocks

Source: https://context7.com/lucidrains/x-mlps-pytorch/llms.txt

Implement transformer feedforward blocks using the Feedforwards class. It includes RMSNorm, residual connections, and configurable expansion factors. Supports input/output projection and list/tuple inputs.

```python
import torch
from x_mlps_pytorch import Feedforwards

# 4 feedforward blocks with dim=256, expansion_factor=4
ff_stack = Feedforwards(
    dim=256,
    depth=4,
    expansion_factor=4.,  # hidden dim = 256 * 4 = 1024
    final_norm=True
)

# With input/output projection
ff_with_proj = Feedforwards(
    dim=512,
    depth=6,
    dim_in=128,   # project 128 -> 512
    dim_out=64,   # project 512 -> 64
    activation=torch.nn.GELU(),
    norm_after_activation=True  # additional norm after activation
)

# Forward pass (includes residual connections)
x = torch.randn(32, 128)
output = ff_with_proj(x)  # (32, 64)

# Also supports list/tuple inputs
obs = torch.randn(32, 64)
goal = torch.randn(32, 64)
output = ff_with_proj([obs, goal])  # concatenates then processes
```

--------------------------------

### Hyperspherical Normalized Feedforward Networks

Source: https://context7.com/lucidrains/x-mlps-pytorch/llms.txt

Utilize nFeedforwards for hyperspherical normalization as seen in the nGPT paper. This class uses L2-normalized activations and SLERP-style residual updates for stable training. Options include magnitude preservation and manual weight normalization.

```python
import torch
from x_mlps_pytorch import nFeedforwards

# Basic nFeedforwards
nff = nFeedforwards(
    dim=512,
    depth=4,
    ff_expand_factor=4.
)

# With input/output projections and magnitude preservation
nff_full = nFeedforwards(
    dim=512,
    depth=4,
    dim_in=128,
    dim_out=128,
    input_preserve_magnitude=True,  # preserve input magnitude info (SimBa v2 technique)
    constant_shift=3.,  # constant for magnitude preservation
    alpha_init=0.1,  # initial residual mixing coefficient
    manual_norm_weights=False  # use parametrization for normalized weights
)

# Forward pass
x = torch.randn(2, 128)
output = nff_full(x)  # (2, 128)

# Manually normalize weights (if using manual_norm_weights=True)
# nff_full.norm_weights_()
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.