Limitations

Let's see this mistake in action:

using SparseConnectivityTracer
detector = TracerSparsityDetector()

relu_bad(x::AbstractFloat) = max(zero(x), x)

f_bad(xs) = sum(relu_bad, xs)

Since tracers and dual numbers are Real numbers and not AbstractFloats, relu_bad throws a MethodError:

julia> xs = [1.0, -2.0, 3.0];
julia> f_bad(xs)4.0
julia> jacobian_sparsity(f_bad, xs, detector)ERROR: MethodError: no method matching relu_bad(::SparseConnectivityTracer.GradientTracer{SparseConnectivityTracer.IndexSetGradientPattern{Int64, BitSet}})
The function `relu_bad` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  relu_bad(::AbstractFloat)
   @ Main limitations.md:29

This is easily fixed by loosening type restrictions or adding an additional methods on Real:

relu_good(x) = max(zero(x), x)

f_good(xs) = sum(relu_good, xs)

julia> jacobian_sparsity(f_good, xs, detector)1×3 SparseArrays.SparseMatrixCSC{Bool, Int64} with 3 stored entries:
 1  1  1

Let us motivate the design decision above by a simple example function:

function f(x)
    if x[1] > x[2]
        return x[1]
    else
        return x[2]
    end
end

The desired global Jacobian sparsity pattern over the entire input domain $x \in \mathbb{R}^2$ is [1 1]. Two local sparsity patterns are possible: [1 0] for $\{x | x_1 > x_2\}$, [0 1] for $\{x | x_1 \le x_2\}$.

The local sparsity patterns of TracerLocalSparsityDetector are easy to compute using operator overloading by using dual numbers which contain primal values on which we can evaluate comparisons like >:

using SparseConnectivityTracer

jacobian_sparsity(f, [2, 1], TracerLocalSparsityDetector())

jacobian_sparsity(f, [1, 2], TracerLocalSparsityDetector())

1×2 SparseArrays.SparseMatrixCSC{Bool, Int64} with 1 stored entry:
 ⋅  1

The global sparsity pattern is impossible to compute when code branches with an if-else condition, since we can only ever hit one branch during run-time. If we made comparisons like > return true or false, we'd get the local patterns [1 0] and [0 1] respectively. But SCT's policy is to guarantee conservative sparsity patterns, which means that "false positives" (ones) are acceptable, but "false negatives" (zeros) are not. In my our opinion, the right thing to do here is to throw an error:

julia> jacobian_sparsity(f, [1, 2], TracerSparsityDetector())ERROR: TypeError: non-boolean (SparseConnectivityTracer.GradientTracer{SparseConnectivityTracer.IndexSetGradientPattern{Int64, BitSet}}) used in boolean context

In some cases, we can work around this by using ifelse. Since ifelse is a method, it can evaluate "both branches" and take a conservative union of both resulting sparsity patterns:

f(x) = ifelse(x[1] > x[2], x[1], x[2])

jacobian_sparsity(f, [1, 2], TracerSparsityDetector())

1×2 SparseArrays.SparseMatrixCSC{Bool, Int64} with 2 stored entries:
 1  1

As motivated in the section above, global sparsity detection isn't allowed to hit any branching code. While SCT's overloads try to avoid branches by throwing errors, they can in some cases be entered by stateful functions. Let's look at a function whose output doesn't only depend on the input x, but also on an internal state (in this case a random number):

julia> using SparseConnectivityTracer
julia> f(x) = randn() > 0.5 ? x[1] : x[2];
julia> jacobian_sparsity(f, [1, 2], TracerSparsityDetector())1×2 SparseArrays.SparseMatrixCSC{Bool, Int64} with 1 stored entry:
 ⋅  1

SCT cannot return the correct global sparsity pattern [1 1] for this code. This issue can be circumvented by adding an overload on f that returns a conservative pattern.

Problems can also arise when statefulness is introduced via mutable caches. In this example, we assume a stateful function f whose output not only depends on the input x, but also on a mutable array defined outside of the function call:

julia> using SparseConnectivityTracer, SparseArrays
julia> A_cache = sparse([2 0; 0 3])2×2 SparseArrays.SparseMatrixCSC{Int64, Int64} with 2 stored entries:
 2  ⋅
 ⋅  3
julia> f(x) = A_cache * x;
julia> pattern1 = jacobian_sparsity(f, [1, 2], TracerSparsityDetector())2×2 SparseArrays.SparseMatrixCSC{Bool, Int64} with 2 stored entries:
 1  ⋅
 ⋅  1

Invalidation of global sparsity patterns

While this sparsity pattern is correct at the time of detection, it can be invalidated by mutating the array A_cache:

julia> A_cache[1, 2] = 4;
julia> A_cache2×2 SparseArrays.SparseMatrixCSC{Int64, Int64} with 3 stored entries:
 2  4
 ⋅  3
julia> pattern2 = jacobian_sparsity(f, [1, 2], TracerSparsityDetector())2×2 SparseArrays.SparseMatrixCSC{Bool, Int64} with 3 stored entries:
 1  1
 ⋅  1

With the wisdom of hindsight, pattern1 can therefore be seen as "non-conservative".

Exception: Dense arrays

For dense caches of type Array (including Matrix and Vector), SCT tries to circumvent this issue by returning a conservative sparsity pattern (note that this behavior could be changed in a future breaking release):

julia> A_cache = [2 0; 0 3]2×2 Matrix{Int64}:
 2  0
 0  3
julia> pattern1 = jacobian_sparsity(f, [1, 2], TracerSparsityDetector())2×2 SparseArrays.SparseMatrixCSC{Bool, Int64} with 4 stored entries:
 1  1
 1  1

However, such guarantees can't be made for arbitrary cache types (like the SparseMatrixCSC above).