Docs / tutorial on sparsity?

[gdalle]

Related to SciML/ADTypes.jl#127 (comment), would it be possible to add a page about sparsity detection / exploitation to the Reactant docs? It gets mentioned regularly but I'm not sure where to start. For instance, what would be the equivalent of the following workflow with Reactant? How do I get a sparse matrix as the output, and what type does it have?

using ForwardDiff
using DifferentiationInterface
using SparseConnectivityTracer, SparseMatrixColorings

f(x) = map(sin, x) + map(cos, reverse(x))
backend = AutoSparse(
    AutoForwardDiff();
    sparsity_detector = TracerSparsityDetector(),
    coloring_algorithm = GreedyColoringAlgorithm()
);
x = rand(10)
prep = prepare_jacobian(f, backend, x);
jacobian(f, prep, backend, x)

[wsmoses]

following up from slack see https://c.wsmoses.com/presentations/reactant.pdf , slides 16-18

for a specific julia example try this:

julia> a = Reactant.to_rarray(ones(100000));

julia> z = Reactant.to_rarray(ones(100000));

julia> function f(a,z)
           b = a .* z
           c = b .+ 4
           return c[1:11]
       end
f (generic function with 1 method)

julia> @code_hlo f(a,z)
module @reactant_f attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<100000xf64> {enzymexla.memory_effects = []}, %arg1: tensor<100000xf64> {enzymexla.memory_effects = []}) -> tensor<11xf64> attributes {enzymexla.memory_effects = []} {
    %cst = stablehlo.constant dense<4.000000e+00> : tensor<11xf64>
    %0 = stablehlo.slice %arg0 [0:11] : (tensor<100000xf64>) -> tensor<11xf64>
    %1 = stablehlo.slice %arg1 [0:11] : (tensor<100000xf64>) -> tensor<11xf64>
    %2 = stablehlo.multiply %0, %1 : tensor<11xf64>
    %3 = stablehlo.add %2, %cst : tensor<11xf64>
    return %3 : tensor<11xf64>
  }
}

essentially in spite of the dense algebra happening by default, we can optimize it to sparsely compute/store/etc only what is needed.

[lassepe]

My understanding is that the matrix coloring will do more work than just skipping structural zeros; it will also ensure that the columns are computed in fewer backward passes. Is such a feature out of scope?

[wsmoses]

oh not out of scope at all, just describing the currently implemented stuff @gdalle was referring to!

[wsmoses]

there are a couple of other still open issues for sparsity too, like: #308 and #403 but we may want to clean up/make clear each issue's independent pieces.

In any case I closed this since I think the core question from @gdalle was answered, but also feel free to reopen if not, or if there's another aspect still pending!