Bug: NNlib.batched_mul fails with MissingPrimalError

[bdrhill]

Summary

NNlib.batched_mul fails with MissingPrimalError for iszero when using TracerSparsityDetector.

MWE

using SparseConnectivityTracer
using NNlib

detector = TracerSparsityDetector()
jacobian_sparsity(
    x -> vec(batched_mul(reshape(x[1:8], 2, 2, 2), reshape(x[9:16], 2, 2, 2))),
    rand(16),
    detector
)
# ERROR: MissingPrimalError(iszero, GradientTracer...)

Root cause

batched_mul internally calls iszero to check for special cases, which requires primal values.

Expected behavior

Should compute the sparsity pattern of batched matrix multiplication.

Workaround

Use TracerLocalSparsityDetector (untested).