sjlt 0.1.1

Sparse Johnson-Lindenstrauss Transform with CUDA acceleration

Sparse Johnson-Lindenstrauss Transform CUDA Kernel

This is a simple repository for Sparse Johnson-Lindenstrauss Transform with CUDA acceleration for PyTorch.

Features

• GPU-accelerated sparse random projections
• Supports float32, float64, and bfloat16 data types
• Optimized CUDA kernels for high performance
• Easy integration with PyTorch workflows

Installation

Requirements

• Python >= 3.8
• PyTorch >= 1.9.0 with CUDA support
• CUDA Toolkit (version compatible with your PyTorch installation)
• C++ compiler (GCC 7-11 recommended)

Install from PyPI

pip install sjlt

Install from Source

git clone https://github.com/TRAIS-Lab/sjlt
cd sjlt
pip install -e .

Due to the default pip install -e . isolation build behavior, it might help to use pip install --no-build-isolation -e . when you see something like:

RuntimeError:
The detected CUDA version (11.8) mismatches the version that was used to compile
PyTorch (12.6). Please make sure to use the same CUDA versions.

Quick Start

Our sjlt implementation has the following parameters:

• original_dim: input dimension
• proj_dim: output dimension
• c: sparsity parameter, i.e., non-zeros per column (default: 1)
• threads: CUDA threads per block (default: 1024)
• fixed_blocks: CUDA blocks to use (default: 84)

We note that the input is supposed to have batch_dim, i.e., input.shape() should be (batch_size, original_dim) and output.shape() will be (batch_size, proj_dim).

The following is a simple snippet of using our SJLT CUDA kernel:

import torch
from sjlt import SJLTProjection

# Create projection: 1024 -> 128 dimensions with sparsity 4
proj = SJLTProjection(original_dim=1024, proj_dim=128, c=4)

# Project some data
x = torch.randn(100, 1024, device='cuda')
y = proj(x)  # Shape: [100, 128]

Troubleshooting

If installation fails:

1. Ensure CUDA toolkit is installed and nvcc is in PATH
2. Check PyTorch CUDA compatibility: python -c "import torch; print(torch.cuda.is_available())"
3. Try reinstalling: pip install sjlt --no-cache-dir --force-reinstall