johannes-fischer / NNlib.jl

Neural Network primitives with multiple backends

NNlib.jl

This package provides a library of functions useful for neural networks, such as softmax, sigmoid, batched multiplication, convolutions and pooling. Many of these are used by Flux.jl, which loads this package, but they may be used independently.

For use with automatic differentiation, this package defines gradients using ChainRules.jl. These will be seen by various packages including Zygote.jl.

GPU support is provided as package extensions (see the ext/ folder). In order to load the extensions, use the imports

using NNlib, CUDA, cuDNN

for CUDA support, or

using NNlib, AMDGPU

for AMDGPU support.