triton-neural 2.0.3

A comprehensive, modular deep learning library featuring state-of-the-art transformer architectures, latent space models, and memory-efficient attention mechanisms.

Triton Neural - Unified Deep Learning API

Version 2.0.3 - Complete neural network framework built on JAX

A comprehensive, modular deep learning library featuring state-of-the-art transformer architectures, latent space models, and memory-efficient attention mechanisms.

Installation

pip install triton-neural

Import Styles

Triton Neural supports two import styles for maximum flexibility:

Direct Import (Simple)

from triton_neural import *

model = Sequential(Linear(784, 128), ReLU())
optimizer = Adam(learning_rate=0.001)

PyTorch-Style Import (Organized)

import triton_neural as tn

# Core components
model = tn.Sequential(tn.Linear(784, 128), tn.ReLU())

# Training utilities
optimizer = tn.train.Adam(learning_rate=0.001)
params, history = tn.train.fit(model, params, optimizer, ...)

# Utilities
tn.util.print_model_summary(model, params, (784,))
tn.util.plot_history(history)

# Transformers
transformer = tn.transformer.PremadeTransformer(num_layers=6, embed_dim=512)

Both styles work identically - choose what fits your workflow!

Features

Core Components

• Basic Layers: Linear, Conv2D, BatchNorm, Dropout, Flatten, Reshape
• Activations: ReLU, GELU, Sigmoid, Tanh, Softmax, LeakyReLU
• Container: Sequential for easy model building

Advanced Transformers

• 6 Attention Types:
  • self - Standard multi-head attention
  • masked - Causal attention for GPT-style models
  • cross - Encoder-decoder attention
  • sparse - O(n√n) for long sequences
  • flash - O(n) memory for very long contexts
  • rope - Rotary Position Embeddings (modern LLMs)
• Normalization: LayerNorm, RMSNorm
• Position Encodings: Sinusoidal, Learned, RoPE
• Complete Models: PremadeTransformer, PremadeTransformerDecoder

Latent Spaces

• VAE: Variational Autoencoder
• CVAE: Conditional VAE
• Utilities: Interpolation, attribute vectors, latent manipulation

Training & Optimization

• Optimizers: SGD, Adam, RMSprop
• Loss Functions: MSE, CrossEntropy, VAE losses
• Training Loop: Complete fit() with validation
• Utilities: Model summary, plotting, save/load

Quick Start

Basic Neural Network

import triton_neural as tn
from jax import random

# Create a simple MLP
model = tn.Sequential(
    tn.Linear(784, 256),
    tn.ReLU(),
    tn.Dropout(0.2),
    tn.Linear(256, 10),
    tn.Softmax()
)

# Initialize
rng = random.PRNGKey(0)
params = model.init(rng, (784,))

# Train
optimizer = tn.train.Adam(learning_rate=0.001)
optimizer.init(params)

params, history = tn.train.fit(
    model, params, optimizer,
    train_data=(x_train, y_train),
    epochs=10,
    loss_fn=tn.cross_entropy_loss
)

Modern Transformer (Flash Attention)

import triton_neural as tn
from jax import random

# Memory-efficient transformer for long sequences
transformer = tn.transformer.PremadeTransformer(
    num_layers=12,
    embed_dim=768,
    num_heads=12,
    attention_type='flash',  # O(n) memory!
    max_len=2048
)

rng = random.PRNGKey(0)
params = transformer.init(rng, (2048, 768))

# Process long sequence
x = random.normal(rng, (1, 1024, 768))
output = transformer(x, params, rng, training=False)

GPT-style Language Model

import triton_neural as tn

# Autoregressive decoder with RoPE
gpt = tn.transformer.PremadeTransformer(
    num_layers=12,
    embed_dim=768,
    num_heads=12,
    attention_type='rope',  # Modern position encoding
    max_len=2048,
    use_learned_pos=False
)

params = gpt.init(rng, (2048, 768))

Variational Autoencoder

import triton_neural as tn

# VAE for latent space learning
vae = tn.VAE(
    input_dim=784,
    latent_dim=32,
    encoder_hidden=[256, 128],
    decoder_hidden=[128, 256]
)

params = vae.init(rng, (784,))

# Encode, sample, decode
reconstruction, mu, logvar = vae(x, params, rng, training=True)

# Compute VAE loss
loss = tn.vae_loss(reconstruction, x, mu, logvar, beta=1.0)

# Latent space utilities
latent_space = tn.LatentSpace(vae)
latent = latent_space.read(x, params, rng)

# Interpolate in latent space
z1 = latent['sample'][0]
z2 = latent['sample'][1]
interpolated = latent_space.interpolate(z1, z2, steps=10)

Attention Type Selection

When to Use What

Attention Type	Complexity	Memory	Use Case
self	O(n²)	O(n²)	Standard, n < 512
masked	O(n²)	O(n²)	GPT-style, autoregressive
sparse	O(n√n)	O(n√n)	Long docs, 512-2048
flash	O(n²)	O(n)	Memory-limited, 2048+
rope	O(n²)	O(n²)	Modern LLMs, no length limit
cross	O(n²)	O(n²)	Encoder-decoder

Recommendations

import triton_neural as tn

# Sequence < 512 tokens
transformer = tn.transformer.PremadeTransformer(attention_type='self')

# GPT-style generation
gpt = tn.transformer.PremadeTransformer(attention_type='masked')

# Long documents (512-2048)
doc_encoder = tn.transformer.PremadeTransformer(attention_type='sparse')

# Very long sequences (2048+) or limited memory
efficient_model = tn.transformer.PremadeTransformer(attention_type='flash')

# Building modern LLM
llm = tn.transformer.PremadeTransformer(attention_type='rope')

Architecture Examples

BERT-style Encoder

import triton_neural as tn

bert = tn.transformer.PremadeTransformer(
    num_layers=12,
    embed_dim=768,
    num_heads=12,
    ff_dim=3072,
    max_len=512,
    attention_type='self',
    dropout=0.1
)

GPT-style Decoder

import triton_neural as tn

gpt = tn.transformer.PremadeTransformer(
    num_layers=12,
    embed_dim=768,
    num_heads=12,
    ff_dim=3072,
    max_len=1024,
    attention_type='masked',
    use_learned_pos=True,
    dropout=0.1
)

Seq2Seq Translation

import triton_neural as tn

# Encoder
encoder = tn.transformer.PremadeTransformer(
    num_layers=6,
    embed_dim=512,
    num_heads=8
)

# Decoder with cross-attention
decoder = tn.transformer.PremadeTransformerDecoder(
    num_layers=6,
    embed_dim=512,
    num_heads=8,
    use_cross_attention=True
)

# Usage
encoder_out = encoder(source, enc_params, rng)
decoder_out = decoder(target, dec_params, rng, 
                     encoder_output=encoder_out)

Long Document Processing

import triton_neural as tn

# Efficient sparse attention for 2048 tokens
long_encoder = tn.transformer.PremadeTransformer(
    num_layers=8,
    embed_dim=512,
    num_heads=8,
    max_len=4096,
    attention_type='sparse'  # O(n√n) complexity
)

Module Organization

Triton Neural is organized into modules for easy access:

Core Module (triton_neural)

Direct access to layers, activations, and basic components:

import triton_neural as tn

tn.Linear(784, 128)
tn.ReLU()
tn.Sequential(...)
tn.VAE(...)

Train Module (triton_neural.train)

Training utilities, optimizers, and model persistence:

tn.train.Adam(learning_rate=0.001)
tn.train.fit(model, params, optimizer, ...)
tn.train.save_params(params, 'model.pkl')
tn.train.load_params('model.pkl')

Util Module (triton_neural.util)

Visualization and model inspection:

tn.util.print_model_summary(model, params, input_shape)
tn.util.plot_history(history)
tn.util.print_attention_guide()

Transformer Module (triton_neural.transformer)

All transformer components and attention mechanisms:

tn.transformer.PremadeTransformer(...)
tn.transformer.SelfAttention(...)
tn.transformer.FlashAttention(...)
tn.transformer.RoPEAttention(...)

Complete Examples

See triton_neural_examples.py for 13 comprehensive examples:

1. MLP Classifier
2. CNN for Images
3. BERT-style Transformer
4. GPT-style Decoder
5. Sparse Attention
6. Flash Attention
7. RoPE Attention
8. Seq2Seq Translation
9. VAE
10. Latent Space Manipulation
11. Conditional VAE
12. Hybrid Models
13. Complete Training Pipeline

API Reference

Core Layers

import triton_neural as tn

tn.Linear(in_features, out_features)
tn.Conv2D(in_channels, out_channels, kernel_size=3, stride=1, padding='SAME')
tn.BatchNorm(num_features, momentum=0.9, eps=1e-5)
tn.Dropout(rate=0.5)

Activations

tn.ReLU()
tn.GELU()
tn.Sigmoid()
tn.Tanh()
tn.Softmax()
tn.LeakyReLU(negative_slope=0.01)

Attention Mechanisms

tn.transformer.SelfAttention(embed_dim, num_heads=8, dropout=0.0)
tn.transformer.MaskedSelfAttention(embed_dim, num_heads=8, dropout=0.0)
tn.transformer.CrossAttention(embed_dim, num_heads=8, dropout=0.0)
tn.transformer.SparseAttention(embed_dim, num_heads=8, block_size=64, stride=64)
tn.transformer.FlashAttention(embed_dim, num_heads=8, block_size=64)
tn.transformer.RoPEAttention(embed_dim, num_heads=8, max_len=2048)

Complete Models

tn.transformer.PremadeTransformer(
    num_layers, embed_dim, num_heads=8,
    ff_dim=None, max_len=512, dropout=0.1,
    attention_type='self',  # 'self', 'masked', 'sparse', 'flash', 'rope'
    use_learned_pos=False
)

tn.transformer.PremadeTransformerDecoder(
    num_layers, embed_dim, num_heads=8,
    ff_dim=None, max_len=512, dropout=0.1,
    use_cross_attention=True,
    use_learned_pos=False
)

Latent Space Models

tn.VAE(input_dim, latent_dim, encoder_hidden=None, decoder_hidden=None)
tn.ConditionalVAE(input_dim, latent_dim, num_classes, ...)
tn.LatentSpace(vae)  # Utility class for manipulation

Optimizers

tn.train.SGD(learning_rate=0.01, momentum=0.0)
tn.train.Adam(learning_rate=0.001, beta1=0.9, beta2=0.999)
tn.train.RMSprop(learning_rate=0.01, decay=0.9)

Loss Functions

tn.mse_loss(predictions, targets)
tn.cross_entropy_loss(logits, labels)
tn.binary_cross_entropy_loss(predictions, targets)
tn.vae_loss(reconstruction, x, mu, logvar, beta=1.0)
tn.vae_mse_loss(reconstruction, x, mu, logvar, beta=1.0)

Training

tn.train.fit(model, params, optimizer, train_data, epochs, loss_fn,
             val_data=None, batch_size=32, rng=None, verbose=True)

tn.train.train_step(model, params, optimizer, x, y, loss_fn, rng=None)
tn.train.eval_step(model, params, x, y, loss_fn)

Utilities

tn.train.accuracy(predictions, targets)
tn.util.plot_history(history, metric='loss')
tn.util.print_model_summary(model, params, input_shape)
tn.train.save_params(params, filepath)
tn.train.load_params(filepath)

Key Benefits

Modularity

• Small, composable functions
• Mix and match any components
• Clear, explicit data flow
• PyTorch-style module organization

Performance

• Built on JAX (JIT compilation, GPU support)
• O(n) memory with Flash Attention
• Efficient sparse attention for long sequences

Modern Features

• 6 attention types
• RoPE for state-of-the-art position encoding
• Flash Attention for memory efficiency
• Complete VAE support with utilities

Educational

• Clean, readable code
• Comprehensive examples
• Well-documented API
• Two import styles for flexibility

Advanced Topics

Memory Optimization

import triton_neural as tn

# Use Flash Attention for long sequences
model = tn.transformer.PremadeTransformer(
    attention_type='flash',  # O(n) memory
    max_len=2048
)

# Use sparse attention for efficiency
model = tn.transformer.PremadeTransformer(
    attention_type='sparse',  # O(n√n)
    max_len=2048
)

Latent Space Arithmetic

import triton_neural as tn

latent_space = tn.LatentSpace(vae)

# Interpolate between points
z_interp = latent_space.interpolate(z1, z2, steps=10)

# Attribute vectors (e.g., "smile" direction)
smile_vec = latent_space.get_attribute_vector(
    z_positive=smiling_faces,
    z_negative=neutral_faces
)

# Apply attribute
z_smiling = latent_space.apply_attribute(z, smile_vec, strength=2.0)

Hybrid Models

import triton_neural as tn

# Combine different attention types
class HybridModel:
    def __init__(self):
        # Early layers: Flash (memory efficient)
        self.flash_layers = tn.transformer.PremadeTransformer(
            num_layers=6,
            attention_type='flash'
        )
        
        # Later layers: RoPE (better positions)
        self.rope_layers = tn.transformer.PremadeTransformer(
            num_layers=6,
            attention_type='rope'
        )
    
    def __call__(self, x, params, rng):
        x = self.flash_layers(x, params['flash'], rng)
        x = self.rope_layers(x, params['rope'], rng)
        return x

Performance Characteristics

Component	Time	Space	Notes
Linear	O(nd)	O(nd)	Standard
Conv2D	O(n·k²·c)	O(n·c)	k=kernel, c=channels
SelfAttention	O(n²d)	O(n²)	Standard
SparseAttention	O(n√n·d)	O(n√n)	Efficient
FlashAttention	O(n²d)	O(n)	Memory-efficient
RoPEAttention	O(n²d)	O(n²)	Modern LLMs
Full Transformer	O(L·n²d)	O(L·n²)	L=layers

Development

Requirements

• Python 3.8+
• JAX 0.6.0+
• NumPy

Structure

triton_neural/
├── __init__.py          # Core layers, activations, VAE
├── train.py             # Optimizers, training loops
├── util.py              # Visualization, utilities
└── transformer.py       # Attention mechanisms, transformers

Citation

If you use Triton Neural in your research, please cite:

@software{triton_neural,
  title={Triton Neural: Unified Deep Learning API},
  author={Built with JAX},
  year={2026},
  version={2.0.2}
}

Contributing

Contributions welcome! This is a modular, educational framework.

License

MIT License - Free to use in research and production.

Acknowledgments

• Built on JAX by Google
• Inspired by PyTorch, Flax, and modern transformer architectures
• Flash Attention algorithm from Dao et al.
• RoPE from Su et al. (used in LLaMA, PaLM)

What Makes Triton Neural Special

1. True Modularity - Every component is independent
2. PyTorch-Style Organization - tn.train.Adam, tn.util.plot_history, etc.
3. 6 Attention Types - More than most frameworks
4. Memory-Efficient - Flash Attention with O(n) memory
5. Modern LLM Support - RoPE, modern architectures
6. Complete VAE Suite - With latent space utilities
7. Educational - Clean, readable, well-documented
8. Production-Ready - Full training pipelines, save/load
9. Flexible - Two import styles to match your preference

Learn More

• See triton_neural_examples.py for 13 working examples
• Check inline documentation for detailed API info
• Experiment with different attention types
• Build hybrid models combining multiple approaches

---

Happy Deep Learning with Triton Neural!

A complete, modular deep learning framework for modern neural networks.