first-kernel

A custom kernel for PyTorch.

Installation

pip install drbh/first-kernel

Usage

import torch
from first_kernel import first_kernel

# Create input tensor
x = torch.randn(1024, 1024, device="cuda")

# Run kernel
result = first_kernel(x)

Development

Building

nix develop
nix run .#build-and-copy

Testing

nix develop .#test
pytest tests/

Test as a kernels user

uv run example.py

License

Apache 2.0