Warp-NN: CUDA Graphable Neural Networks for NVIDIA Warp#
Warp-NN is a Python library for building and training neural networks using NVIDIA Warp. It enables end-to-end GPU-accelerated neural network implementation for Warp-based applications, without the need for third-party libraries such as PyTorch or JAX.
Note
Warp-NN is designed with a focus on clean, simple, and readable neural network-related code. Although the library strives to offer high-performance implementations, they may not always outperform other highly optimized solutions for NVIDIA GPUs that use cuBLAS / cuDNN.
Main features:
Clean, simple, and readable code.
Configurable CUDA kernels.
CUDA Graphable implementation.
Quickstart#
The easiest way to install Warp-NN is from PyPI. Refer to the Installation section for more details.
pip install warp-nn