attorch provides a curated subset of PyTorch's neural network modules implemented in Python using OpenAI's Triton. It targets developers seeking a more hackable, efficient, and readable alternative to pure PyTorch for custom deep learning operations, particularly those who find writing raw CUDA kernels challenging. The library supports both forward and backward passes, enabling its use during model training.

How It Works

attorch leverages OpenAI's Triton, a language and compiler for writing high-performance kernels that bridge the gap between Python and CUDA. By implementing core neural network operations (like convolutions, attention, and various activation functions) in Triton, attorch aims to achieve performance improvements over standard PyTorch while maintaining a Pythonic interface. This approach allows for easier customization and potential for operation fusion, which can further boost efficiency.

Quick Start & Requirements

• Install via pip: pip install torch==2.4.0 triton==3.0.0 and clone the repository.
• Dependencies: torch==2.4.0, triton==3.0.0.
• Official Docs: https://github.com/BobMcDear/attorch

Highlighted Details

• Implements a broad range of layers including Conv1d/2d, MultiheadAttention, various activations, BatchNorm, LayerNorm, and common loss functions.
• Supports Automatic Mixed Precision (AMP).
• Includes a PyTorch fallback mechanism for layers not yet implemented in attorch.
• Provides a attorch.math module for custom kernel development and operation fusion.

Maintenance & Community

• The project is maintained by BobMcDear.
• Further community engagement details (e.g., Discord, Slack) are not specified in the README.

Licensing & Compatibility

• The README does not explicitly state a license. This requires clarification for commercial use or closed-source linking.

Limitations & Caveats

• The implementation is a subset of PyTorch's modules; some layers may not be available and will fall back to PyTorch.
• Performance for inference may not match dedicated inference engines.
• Some tests against PyTorch counterparts might fail due to numerical precision differences.