BatchSizeSettings(B, N) is type-unstable even for a fixed batch size, which blocks forward jacobian under juliac --trim #1088

BatchSizeSettings(B::Int, N::Int) builds two of its type parameters from the runtime input length:

function BatchSizeSettings(B::Int, N::Int)
    # type-unstable, unless `B` and `N` are constant-propagated
    singlebatch = B == N
    aligned = (B == N == 0) || (N % B == 0)
    return BatchSizeSettings{B, singlebatch, aligned}(N)
end

This holds even when the backend has a fixed chunk size (AutoForwardDiff(; chunksize = 4)), so B is a compile-time constant. N = length(x) is never a constant, so pick_batchsize and every prep that stores the settings (PushforwardJacobianPrep, PullbackJacobianPrep, HVPGradientHessianPrep) return a non-concrete type. The jacobian call that follows is a dynamic dispatch.

In the JIT that only costs one dispatch per call. Under juliac --trim=safe it's a hard error, because the verifier can't resolve the call. I hit this while adding a Piste backend (PR to follow): prepared gradient, derivative, pushforward, pullback and reverse-mode jacobian all trim cleanly through DI, but forward-mode jacobian doesn't. Nothing in the error is Piste-specific, so I'd expect the same failure for any forward backend that goes through pick_batchsize.

module DIJacCLI

using ADTypes: AutoPiste
using DifferentiationInterface: prepare_jacobian, jacobian!
using Piste: Piste

h(x) = sin.(x) .* x[1]

function (@main)(args::Vector{String})::Cint
    x = Float64[parse(Float64, a) for a in args]
    n = length(x)
    fwd = AutoPiste(; chunksize = 2)
    J = zeros(n, n)
    p = prepare_jacobian(h, fwd, x)
    jacobian!(h, J, p, fwd, x)
    for v in J
        println(Core.stdout, v)
    end
    return 0
end

end

Verifier error #1: unresolved call from statement DIJacCLI.jacobian!(DIJacCLI.h, %new()::Matrix{Float64}, DifferentiationInterface.prepare_jacobian_nokwarg(...)...
  [1] main(args::Vector{String}) @ DIJacCLI src/DIJacCLI.jl:18

Verifier error #2: unresolved call from statement DifferentiationInterface._prepare_jacobian_aux(Val{true}(), PushforwardFast(), ...
  [1] prepare_jacobian_nokwarg(::Val{true}, ::typeof(h), ::AutoPiste{ForwardMode, 2}, ::Vector{Float64})
      @ DifferentiationInterface src/first_order/jacobian.jl:180

Verifier error #3: unresolved call from statement (Core.apply_type(DifferentiationInterface.BatchSizeSettings, B::Int64, (B::Int64 === N::Int64)::Bool, φ ()::Bool) ...

Possible fix

The type parameters can only take three combinations, so branching explicitly gives inference a small Union that it can split:

function BatchSizeSettings(::Val{B}, N::Int) where {B}
    if B == N
        return BatchSizeSettings{B, true, true}(N)
    elseif B == 0 || N % B == 0
        return BatchSizeSettings{B, false, true}(N)
    else
        return BatchSizeSettings{B, false, false}(N)
    end
end

Backends with a statically known batch size would call BatchSizeSettings(Val(chunksize), N). BatchSizeSettings(B::Int, N) could stay as it is for the adaptive case.

I tried this locally: I added the method above and changed the Piste extension's fixed-chunksize pick_batchsize to call it. The reproducer then builds with 0 verifier errors, and the trimmed binary's Jacobian matches the JIT exactly for N = 2 (single batch), N = 4 (aligned) and N = 3 (unaligned last batch). The other places that build BatchSizeSettings(chunksize, N) from a fixed chunk size (the ForwardDiff and StaticArrays extensions, and AutoSimpleFiniteDiff) would need the same one-line change. I haven't checked whether the batched hessian paths need more.

Happy to send a PR if you're open to this or something like it.