using SparseArrays
using Zygote
using Enzyme


function test_sparse(M, v)
    return sum(abs2, M*v)
end


M = sprand(256, 128, 0.1)
v = rand(128)
dv = zero(v)

f = let M=M
    x->test_sparse(M, x)
end


@benchmark autodiff(Enzyme.Reverse, $test_sparse, $Active, $(Const(M)), Duplicated($v, fill!($dv, 0.0)))
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (min … max):  38.669 μs … 879.459 μs  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     70.779 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   70.571 μs ±  13.251 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

                          ██▂▂                                  
  ▂▁▂▁▁▁▁▁▂▁▂▁▂▃▇▅▄▃▃▃▃▃▄▅████▅▄▃▃▃▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂ ▃
  38.7 μs         Histogram: frequency by time          115 μs <

 Memory estimate: 4.30 KiB, allocs estimate: 3.

@benchmark Zygote.gradient($f, $v)
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (min … max):   9.630 μs … 103.575 ms  ┊ GC (min … max):  0.00% … 99.83%
 Time  (median):     14.730 μs               ┊ GC (median):     0.00%
 Time  (mean ± σ):   57.414 μs ±   1.050 ms  ┊ GC (mean ± σ):  30.57% ±  6.75%

    █                                                           
  ▃███▃▂▂▂▂▂▂▂▂▂▂▂▂▂▂▁▁▁▂▁▂▂▁▂▂▁▂▁▁▁▁▁▁▂▂▂▂▄▅▅▄▄▄▃▃▃▃▂▂▂▂▂▂▂▂▂ ▂
  9.63 μs         Histogram: frequency by time          103 μs <

 Memory estimate: 290.55 KiB, allocs estimate: 26.