burn_jit::kernel

Module conv

Source
Expand description

Convolution kernels

Enums§

Conv2dStrategy
The strategy to be used when launching a convolution kernel.
ConvTranspose2dStrategy
The strategy to be used when launching a conv_transpose kernel.

Functions§

conv2d
Perform a 2D convolution with the given strategy
conv_transpose2d
Perform a 2D convolution with the given strategy
nchw_to_nhwc
Efficiently transpose an NCHW tensor to NHWC for use in kernels that prefer NHWC for performance. Faster than into_contiguous, but specialized only for this specific permutation.