Document GPU support (#145)

Author: adrhill

docs/src/literate/example.jl

@@ -133,3 +133,36 @@ heatmap(expl)
 
 # For more information on heatmapping batches,
 # refer to the [heatmapping documentation](@ref docs-heatmapping-batches).
+
+# ## [GPU support](@id gpu-docs)
+# All analyzers support GPU backends,
+# building on top of [Flux.jl's GPU support](https://fluxml.ai/Flux.jl/stable/gpu/).
+# Using a GPU only requires moving the input array and model weights to the GPU.
+#
+# For example, using [CUDA.jl](https://github.com/JuliaGPU/CUDA.jl):
+
+# ```julia
+# using CUDA, cuDNN
+# using Flux
+# using ExplainableAI
+#
+# # move input array and model weights to GPU
+# input = input |> gpu # or gpu(input)
+# model = model |> gpu # or gpu(model)
+#
+# # analyzers don't require calling `gpu`
+# analyzer = LRP(model)
+#
+# # explanations are computed on the GPU
+# expl = analyze(input, analyzer)
+# ```
+
+# Some operations, like saving, require moving explanations back to the CPU.
+# This can be done using Flux's `cpu` function:
+
+# ```julia
+# val = expl.val |> cpu # or cpu(expl.val)
+#
+# using BSON
+# BSON.@save "explanation.bson" val
+# ```

docs/src/literate/lrp/basics.jl

@@ -118,18 +118,20 @@ expl = analyze(input, analyzer; layerwise_relevances=true)
 expl.extras.layerwise_relevances
 
 # ## [Performance tips](@id docs-lrp-performance)
+# ### Using LRP with a GPU
+# Like all other analyzers, LRP can be used on GPUs.
+# Follow the instructions on [*GPU support*](@ref gpu-docs).
+#
 # ### Using LRP without a GPU
-# Since ExplainableAI.jl's LRP implementation makes use of
-# [Tullio.jl](https://github.com/mcabbott/Tullio.jl),
-# analysis can be accelerated by loading either
-# - a package from the [JuliaGPU](https://juliagpu.org) ecosystem,
-#   e.g. [CUDA.jl](https://github.com/JuliaGPU/CUDA.jl), if a GPU is available
-# - [LoopVectorization.jl](https://github.com/JuliaSIMD/LoopVectorization.jl)
-#   if only a CPU is available.
+# Using Julia's package extension mechanism,
+# ExplainableAI.jl's LRP implementation can optionally make use of
+# [Tullio.jl](https://github.com/mcabbott/Tullio.jl) and
+# [LoopVectorization.jl](https://github.com/JuliaSIMD/LoopVectorization.jl)
+# for faster LRP rules on dense layers.
 #
-# This only requires loading the LoopVectorization.jl package before ExplainableAI.jl:
+# This only requires loading the packages before loading ExplainableAI.jl:
 # ```julia
-# using LoopVectorization
+# using LoopVectorization, Tullio
 # using ExplainableAI
 # ```
 #