Proposed changes

This PR implements weight normalization for MLX, addressing issue #1888. Weight normalization is a reparameterization technique that decouples the magnitude of a weight tensor from its direction, making optimization more efficient by improving the conditioning of the optimization problem. Is particularly important for audio processing, among other applications.

Key Features

• Core C++ implementation of mx.weight_norm with optimized paths for different dimensions
• Python module weight_norm.py with user-friendly API and layer wrappers
• Proper handling of MLX's channel ordering differences from PyTorch
• Workaround for the linalg::norm 2-axes limitation
• Convenience classes for common layer types (Linear, Conv1d, Conv2d)
• Comprehensive test suite validating mathematical properties and cross-framework compatibility

Implementation Details

Core C++ Implementation

The core weight_norm operation is implemented with three different paths based on the number of axes to normalize over:

1. Direct path for 1-2 axes using optimized linalg::norm kernels
2. Reshape-based approach for >2 axes, which:
   • Identifies dimensions to keep vs. normalize
   • Handles special cases: normalizing all dims, keeping one dim, keeping multiple dims
   • Reshapes appropriately to leverage the optimized 2D norm kernel
   • Reshapes results back for broadcasting

Python Layer

The Python implementation:

• Provides a weight_norm function that wraps MLX modules
• Handles dimension ordering differences for different layer types
• Computes initial g parameter as the norm of the original weight
• Overrides the module's forward pass to apply weight normalization on-the-fly
• Includes convenience classes (WeightNormLinear, WeightNormConv1d, WeightNormConv2d)

Testing and Verification

Testing follows a comprehensive two-pronged approach:

1. Mathematical Property Tests

• Verify that the normalized weights have the correct norm (equals g)
• Confirm that the direction of normalized weights matches v
• Validate that changing g correctly scales the weight norms
• Test edge cases like normalizing over all dimensions

2. Cross-Framework Verification

• Compare against PyTorch's weight normalization
• Test both independent implementations and direct weight transfer
• Document expected differences between frameworks and how to achieve exact equivalence

3. Performance Benchmarks

Benchmarked on Apple M3 Max shows MLX outperforms PyTorch MPS:

• Linear layers (1 axis): 4.90x-5.26x speedup
• Conv1d layers (2 axes): 1.46x-2.05x speedup
• Conv2d layers (3 axes): 1.50x-1.76x speedup

Usage Examples

Core API

import mlx.core as mx

# Create parameters
v = mx.random.normal((64, 3, 3))  # Direction tensor
g = mx.random.normal((64, 1, 1))  # Magnitude tensor

# Apply weight normalization
w = mx.weight_norm(v, g, axes=[1, 2])

Module API

import mlx.nn as nn
from mlx.nn.layers.weight_norm import weight_norm

# Apply to existing layer
linear = nn.Linear(10, 20)
linear_wn = weight_norm(linear)

# Use convenience class
conv1d_wn = nn.WeightNormConv1d(16, 32, kernel_size=3)

Resolves #1888.

Checklist

• [ X] I have read the CONTRIBUTING document
• [X ] I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• [X ] I have added tests that prove my fix is effective or that my feature works
• [ X] I have updated the necessary documentation (if needed)