Fix keyword argument add_batch_dim (#8)

Author: adrhill

src/XAIBase.jl

@@ -3,6 +3,9 @@ module XAIBase
 using TextHeatmaps
 using VisionHeatmaps
 
+include("compat.jl")
+include("utils.jl")
+
 # Abstract super type of all XAI methods.
 # Is expected that all methods are callable types that return an `Explanation`:
 #

src/analyze.jl

@@ -1,7 +1,7 @@
 
 const BATCHDIM_MISSING = ArgumentError(
     """The input is a 1D vector and therefore missing the required batch dimension.
-    Call `analyze` with the keyword argument `add_batch_dim=false`."""
+    Call `analyze` with the keyword argument `add_batch_dim=true`."""
 )
 
 """

src/compat.jl

@@ -0,0 +1,18 @@
+# https://github.com/JuliaLang/julia/pull/39794
+if VERSION < v"1.7.0-DEV.793"
+    export Returns
+
+    struct Returns{V} <: Function
+        value::V
+        Returns{V}(value) where {V} = new{V}(value)
+        Returns(value) = new{Core.Typeof(value)}(value)
+    end
+
+    (obj::Returns)(args...; kw...) = obj.value
+    function Base.show(io::IO, obj::Returns)
+        show(io, typeof(obj))
+        print(io, "(")
+        show(io, obj.value)
+        return print(io, ")")
+    end
+end

src/utils.jl

@@ -0,0 +1,21 @@
+"""
+    batch_dim_view(A)
+
+Return a view onto the array `A` that contains an extra singleton batch dimension at the end.
+This avoids allocating a new array.
+
+## Example
+```juliarepl
+julia> A = [1 2; 3 4]
+2×2 Matrix{Int64}:
+ 1  2
+ 3  4
+
+julia> batch_dim_view(A)
+2×2×1 view(::Array{Int64, 3}, 1:2, 1:2, :) with eltype Int64:
+[:, :, 1] =
+ 1  2
+ 3  4
+```
+"""
+batch_dim_view(A::AbstractArray{T,N}) where {T,N} = view(A, ntuple(Returns(:), N + 1)...)

test/test_api.jl

@@ -20,6 +20,11 @@ expl = analyze(input, analyzer)
 expl = analyzer(input)
 @test expl.val == val
 
+# Max activation + add_batch_dim
+input_vec = [1, 2, 3]
+expl = analyzer(input_vec; add_batch_dim=true)
+@test expl.val == val[:, 1:1]
+
 # Ouput selection
 output_neuron = 2
 val = [2 30; 4 25; 6 20]
@@ -29,3 +34,7 @@ expl = analyze(input, analyzer, output_neuron)
 @test isnothing(expl.extras)
 expl = analyzer(input, output_neuron)
 @test expl.val == val
+
+# Ouput selection + add_batch_dim
+expl = analyzer(input_vec, output_neuron; add_batch_dim=true)
+@test expl.val == val[:, 1:1]