Causal depthwise conv1d in CUDA, with a PyTorch interface
Project description
Causal depthwise conv1d in CUDA with a PyTorch interface
Features:
- Support fp32, fp16, bf16.
- Kernel size 2, 3, 4.
How to use
from causal_conv1d import causal_conv1d_fn
def causal_conv1d_fn(x, weight, bias=None, activation=None):
"""
x: (batch, dim, seqlen)
weight: (dim, width)
bias: (dim,)
activation: either None or "silu" or "swish"
out: (batch, dim, seqlen)
"""
Equivalent to:
import torch.nn.functional as F
F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim)[..., :seqlen]
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Close
Hashes for causal_conv1d-1.2.2.post1.tar.gz
Algorithm | Hash digest | |
---|---|---|
SHA256 | 6ed00a95815b4536c3331c3a5d4bf2ddc058c52c8d9ee51d99f446a3073357be |
|
MD5 | 412d6af1f5ae76153d4eb010d60f13fb |
|
BLAKE2b-256 | 51e8ad8c4be8207a3b46bf7161f2af9be8b95d8bfb7195779e1b9c411c48c1b4 |