Metadata-Version: 2.4
Name: JunGrad
Version: 0.2.0
Summary: A robust N-D autograd library with comprehensive ops, NN layers, and optimizers
Home-page: https://github.com/junkim100/JunGrad
Author: JunGrad Contributors
License: MIT License
        
        Copyright (c) 2024 JunGrad Contributors
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
        
        
Project-URL: Homepage, https://github.com/junkim100/JunGrad
Project-URL: Documentation, https://github.com/junkim100/JunGrad#readme
Project-URL: Issues, https://github.com/junkim100/JunGrad/issues
Project-URL: Source, https://github.com/junkim100/JunGrad
Keywords: autograd,numpy,deep-learning,machine-learning
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy>=1.24.0
Provides-Extra: dev
Requires-Dist: black>=24.0.0; extra == "dev"
Requires-Dist: ruff>=0.1.0; extra == "dev"
Requires-Dist: mypy>=1.0.0; extra == "dev"
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
Requires-Dist: build>=1.2.1; extra == "dev"
Requires-Dist: twine>=4.0.2; extra == "dev"
Provides-Extra: sparse
Requires-Dist: scipy>=1.10.0; extra == "sparse"
Provides-Extra: viz
Requires-Dist: graphviz>=0.20.0; extra == "viz"
Provides-Extra: tutorial
Requires-Dist: datasets>=4.4.0; extra == "tutorial"
Requires-Dist: matplotlib>=3.7.0; extra == "tutorial"
Requires-Dist: scikit-learn>=1.2.0; extra == "tutorial"
Requires-Dist: ipykernel>=6.0.0; extra == "tutorial"
Dynamic: home-page
Dynamic: license-file
Dynamic: requires-python

# JunGrad

A robust N-D autograd library with comprehensive operations, neural network layers, and optimizers.

## Features

- **N-D Tensor** with numpy backend and automatic differentiation
- **Comprehensive Operations**:
  - Elementwise ops (add, sub, mul, div, exp, log, pow, etc.)
  - Reductions (sum, mean, max, min, var, std)
  - Linear algebra (matmul, transpose, reshape, concat, stack)
  - Indexing (gather, scatter_add, slice, take)
- **Neural Network Modules**:
  - `Linear`, `Conv1d`, `Embedding`, `LayerNorm`, `Dropout`
  - `Sequential`, `Stack` containers
  - Parameter and state_dict management
- **Optimizers**:
  - `SGD` (with momentum and Nesterov)
  - `Adam`, `AdamW` (decoupled weight decay)
  - `RMSProp`
  - Gradient clipping utilities
- **Learning Rate Schedulers**:
  - `StepLR`, `ExponentialLR`
  - `CosineLR` (with warmup)
  - `OneCycleLR`
- **Loss Functions**:
  - `cross_entropy` (with label smoothing support)
  - `mse_loss`
  - `bce_with_logits` (numerically stable)
- **Functional API**:
  - Activations: `relu`, `tanh`, `sigmoid`, `gelu`
  - Stable: `softmax`, `log_softmax`, `logsumexp`
- **Utilities**:
  - Gradient checking via finite differences
  - Profiler for timing operations
  - Graphviz export for computation graphs
  - Hooks system

## Installation

```bash
# From PyPI (once published)
pip install jungrad

# Optional extras
pip install "JunGrad[viz]"        # Graphviz rendering support
pip install "JunGrad[tutorial]"   # Dependencies for the quickstart notebook
pip install "JunGrad[sparse]"     # SciPy-backed sparse utilities

# Development installation (editable)
pip install -e .

# Development with extras
pip install -e ".[dev,sparse,viz,tutorial]"
```

## Quick Start

### Basic Tensor Operations

```python
from jungrad import tensor, randn

# Create tensors
a = tensor([1.0, 2.0, 3.0], requires_grad=True)
b = tensor([4.0, 5.0, 6.0], requires_grad=True)

# Operations
c = a + b
d = a * b
e = a ** 2

# Backward pass
e.backward()
print(f"a.grad = {a.grad}")
print(f"b.grad = {b.grad}")
```

### Neural Network

```python
from jungrad import randn, no_grad
from jungrad.nn import Linear, Sequential, ReLU, Dropout
from jungrad.optim import Adam
from jungrad.losses import cross_entropy
from jungrad.sched import CosineLR
from jungrad.optim import clip_grad_norm_

# Create model with dropout
model = Sequential(
    Linear(10, 64),
    ReLU(),
    Dropout(0.3),
    Linear(64, 32),
    ReLU(),
    Dropout(0.3),
    Linear(32, 5),  # 5 classes
)

# Training data
x = randn(32, 10)
y = [0, 1, 2, 3, 4] * 6 + [0, 1]  # class labels

# Forward pass
logits = model(x)
loss = cross_entropy(logits, y)

# Backward and optimize
optimizer = Adam(model.parameters(), lr=0.001)
loss.backward()
clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
optimizer.zero_grad()

# Evaluation mode
model.eval()
with no_grad():
    pred = model(x)
```

### Complete Training Example

```python
from jungrad import randn, tensor, no_grad
from jungrad.nn import Linear, Sequential, ReLU, Dropout
from jungrad.nn.init import kaiming_normal_
from jungrad.losses import cross_entropy
from jungrad.optim import Adam, clip_grad_norm_
from jungrad.sched import CosineLR
import numpy as np

# Model
model = Sequential(
    Linear(10, 64), ReLU(), Dropout(0.3),
    Linear(64, 32), ReLU(), Dropout(0.3),
    Linear(32, 5),
)

# Initialize weights
for module in model.modules():
    if isinstance(module, Linear):
        kaiming_normal_(module.weight)

# Setup training
optimizer = Adam(model.parameters(), lr=0.001)
scheduler = CosineLR(optimizer, T_max=50)

# Training loop
for epoch in range(50):
    x = randn(32, 10)
    y = np.random.randint(0, 5, size=32)

    model.train()
    logits = model(x)
    loss = cross_entropy(logits, y)

    optimizer.zero_grad()
    loss.backward()
    clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()
    scheduler.step()

    # Validation
    if epoch % 10 == 0:
        model.eval()
        with no_grad():
            val_logits = model(randn(16, 10))
            print(f"Epoch {epoch}, Loss: {loss.item():.4f}")
```

### Gradient Checking

```python
from jungrad import randn
from jungrad.ops import matmul
from jungrad.testing import gradcheck

# Verify gradients are correct
a = randn(3, 4, requires_grad=True)
b = randn(4, 5, requires_grad=True)

passed = gradcheck(lambda x, y: matmul(x, y), (a, b), eps=1e-5)
print(f"Gradients verified: {passed}")
```

## Tutorial

See `quickstart.ipynb` for a comprehensive interactive tutorial covering:
- Creating and manipulating N-D tensors
- Automatic differentiation and gradient computation
- Building neural networks with Sequential layers
- Training models with optimizers (SGD, Adam, AdamW, RMSProp)
- Advanced layers (Conv1d, LayerNorm, Embedding, Dropout)
- Activation functions and functional API
- Learning rate schedulers (StepLR, CosineLR, OneCycleLR)
- Gradient clipping for stable training
- Performance optimization with `no_grad()`
- Computation graph visualization with Graphviz
- Performance profiling
- Gradient checking for verification
- **End-to-end classification demo** with real-world data (20 Newsgroups dataset)

## Documentation

### Tensor API

```python
from jungrad import Tensor, tensor, zeros, ones, randn

# Creation
x = tensor([1.0, 2.0, 3.0], requires_grad=True)
x = zeros((3, 4))
x = ones((2, 2))
x = randn(5, 5)

# Methods
x.detach()           # Detach from graph
x.retain_grad()      # Retain gradient for non-leaf
x.item()             # Get scalar value
x.numpy()            # Get numpy array
x.backward()         # Compute gradients
x.zero_grad()        # Zero gradients
```

### Autograd

```python
from jungrad import no_grad, enable_grad

# Context managers
with no_grad():
    # Operations without gradients
    y = x * 2

with enable_grad():
    # Enable gradients
    z = x + y
```

### Operations

```python
from jungrad.ops import add, mul, matmul, sum, mean

# Elementwise
c = add(a, b)
c = mul(a, b)

# Matrix multiplication
out = matmul(x, w)

# Reductions
s = sum(x, axis=0, keepdims=True)
m = mean(x, axis=-1)
```

### Neural Network Layers

```python
from jungrad.nn import Linear, Conv1d, Embedding, LayerNorm, Dropout

# Linear layer
linear = Linear(in_features=10, out_features=5, bias=True)

# Conv1d
conv = Conv1d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)

# Embedding
embed = Embedding(num_embeddings=1000, embedding_dim=128)

# LayerNorm
ln = LayerNorm(normalized_shape=128, eps=1e-5, affine=True)

# Dropout
dropout = Dropout(p=0.5)
```

### Optimizers

```python
from jungrad.optim import SGD, Adam, AdamW, RMSProp, clip_grad_norm_

# SGD
optimizer = SGD(model.parameters(), lr=0.01, momentum=0.9, nesterov=True)

# Adam
optimizer = Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999))

# AdamW (decoupled weight decay)
optimizer = AdamW(model.parameters(), lr=0.001, weight_decay=0.01)

# Gradient clipping
clip_grad_norm_(model.parameters(), max_norm=1.0)
```

### Loss Functions

```python
from jungrad.losses import cross_entropy, mse_loss, bce_with_logits

# MSE loss
loss = mse_loss(pred, target)

# Cross-entropy (supports index and one-hot targets)
loss = cross_entropy(logits, target_indices)
loss = cross_entropy(logits, target_onehot, label_smoothing=0.1)

# BCE with logits
loss = bce_with_logits(logits, target)
```

## Development

### Running Tests

**Run all tests:**
```bash
pytest tests/
# or
python -m pytest tests/
```

**Run with verbose output:**
```bash
pytest tests/ -v
```

**Run a specific test file:**
```bash
pytest tests/test_tensor.py
```

**Run a specific test:**
```bash
pytest tests/test_autograd.py::test_backward_simple
```

**Run tests matching a pattern:**
```bash
pytest tests/ -k "backward"
```

**With coverage report:**
```bash
pytest tests/ --cov=jungrad --cov-report=term-missing
```

See [TESTING.md](TESTING.md) for detailed testing documentation.

### Code Quality

```bash
# Format
black jungrad/

# Lint
ruff check jungrad/

# Type check
mypy jungrad/
```

## Architecture

- **tensor.py**: N-D Tensor class with autograd support
- **autograd.py**: Backward engine with topological sort
- **ops.py**: Low-level primitive operations (200+ operations)
- **functional.py**: High-level differentiable functions (activations, softmax)
- **losses.py**: Loss functions (MSE, cross-entropy, BCE)
- **nn/**: Neural network modules (Module, Parameter, layers, initialization, utils)
- **optim/**: Optimizers (SGD, Adam, AdamW, RMSProp) and gradient clipping
- **sched/**: Learning rate schedulers (StepLR, CosineLR, OneCycleLR, ExponentialLR)
- **testing/**: Gradient checking utilities
- **graphviz.py**: Computation graph visualization
- **profiler.py**: Performance profiling tools
- **hooks.py**: Forward and backward hooks for debugging

## Design Philosophy

- **N-D from the start**: Not scalar-only like micrograd - works with tensors of any shape
- **Separate engine**: Clean separation between tensor and autograd engine
- **Broadcast-aware**: Proper handling of broadcasting in forward and backward passes
- **Numerically stable**: Uses logsumexp for softmax, stable sigmoid, stable BCE with logits
- **Type-safe**: Full type hints throughout the codebase
- **PyTorch-inspired API**: Familiar interface for users coming from PyTorch
- **Educational**: Clear, well-documented code suitable for learning autograd internals

## Real-World Example

The `quickstart.ipynb` tutorial includes a complete end-to-end classification demo using the [20 Newsgroups dataset](https://huggingface.co/datasets/SetFit/20_newsgroups), demonstrating:
- Loading real-world text data
- TF-IDF feature extraction
- Multi-layer neural network with dropout
- Training with learning rate scheduling
- Gradient clipping for stability
- Comprehensive evaluation metrics
- Training curve visualization

This showcases how JunGrad can be used for practical machine learning tasks beyond synthetic examples.

## Changelog

### Version 0.2.0 (2025-01-XX)

**Major Features:**
- **GPU Support via CuPy**: Full GPU acceleration support using CuPy arrays
  - Automatic device detection and conversion between NumPy and CuPy
  - Device-aware operations throughout the library
  - Seamless fallback to CPU when CuPy is not available
  - New `backend` module with utilities: `has_cupy()`, `get_array_module()`, `to_device_array()`, `to_numpy_array()`

**Improvements:**
- All operations (`add`, `sub`, `mul`, `div`, `matmul`, `maximum`, etc.) now handle mixed NumPy/CuPy arrays
- Autograd engine is device-aware and maintains device consistency during backpropagation
- Embedding layer backward pass updated for GPU compatibility
- `asarray` utility preserves CuPy arrays instead of converting to NumPy

**Technical Details:**
- Operations automatically detect array module (NumPy or CuPy) from input tensors
- Device conversion handled transparently during tensor operations
- Gradient accumulation maintains device consistency across the computation graph

### Version 0.1.0 (Initial Release)

- Core N-D tensor implementation with autograd
- Comprehensive operations library (200+ operations)
- Neural network modules (Linear, Conv1d, Embedding, LayerNorm, Dropout)
- Optimizers (SGD, Adam, AdamW, RMSProp)
- Learning rate schedulers (StepLR, CosineLR, OneCycleLR, ExponentialLR)
- Loss functions (MSE, cross-entropy, BCE with logits)
- Functional API with activations
- Gradient checking utilities
- Computation graph visualization
- Performance profiling tools

## License

MIT License
