# Type stability of \`Lux.batched\_jacobian\`
**URL:** https://discourse.julialang.org/t/type-stability-of-lux-batched-jacobian/128526
**Category:** Machine Learning
**Tags:** question, type-stability, autodiff, lux
**Created:** [April 29, 2025, 3:03pm UTC](https://discourse.julialang.org/t/type-stability-of-lux-batched-jacobian/128526 "2025-04-29T15:03:48Z")
**Posts on this page:** 6
**Page:** 1
### Author:  [@NoFishLikeIan](https://discourse.julialang.org/u/NoFishLikeIan)
#### Post date: [April 29, 2025, 3:03pm UTC](https://discourse.julialang.org/t/type-stability-of-lux-batched-jacobian/128526/1 "2025-04-29T15:03:48Z")
How can I make the call to `Lux.batched_jacobian` of a `StatefulLuxLayer` type stable? I am using `Lux v1.12.4`.
Consider this MWE:
```julia
using Lux
using Random; rng = Xoshiro(42)
input, output = 6, 2
model = Chain(Dense(input => input^2, tanh), Dense(input^2 => output));
ps, st = Lux.setup(rng, model);
n = 100
x = rand(rng, Float32, input, n);
f = StatefulLuxLayer{true}(model, ps, st)
@code_warntype f(x) # Type stable, return type `Matrix{Float32}`
const backend = AutoForwardDiff()
@code_warntype batched_jacobian(f, backend, x) # Type unstable, return type `Any`
```
Am I missing something?
Edit: typo in MWE
---
### Author:  [@JADekker](https://discourse.julialang.org/u/JADekker)
#### Post date: [April 29, 2025, 3:53pm UTC](https://discourse.julialang.org/t/type-stability-of-lux-batched-jacobian/128526/2 "2025-04-29T15:53:38Z")
I think you’re missing `m = 6` in the MWE?
---
### Author:  [@NoFishLikeIan](https://discourse.julialang.org/u/NoFishLikeIan)
#### Post date: [April 29, 2025, 3:55pm UTC](https://discourse.julialang.org/t/type-stability-of-lux-batched-jacobian/128526/3 "2025-04-29T15:55:12Z")
Fixed, thanks!
---
### Author:  [@avikpal](https://discourse.julialang.org/u/avikpal)
#### Post date: [April 29, 2025, 8:55pm UTC](https://discourse.julialang.org/t/type-stability-of-lux-batched-jacobian/128526/4 "2025-04-29T20:55:40Z")
> [@NoFishLikeIan](#):
>
> `@code_warntype batched_jacobian(f, backend, x)`
You are missing the chunksize
```julia
julia> @code_warntype batched_jacobian(f, AutoForwardDiff(; chunksize=8), x)
MethodInstance for batched_jacobian(::StatefulLuxLayer{Static.True, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, @NamedTuple{layer_1::@NamedTuple{weight::Matrix{Float32}, bias::Vector{Float32}}, layer_2::@NamedTuple{weight::Matrix{Float32}, bias::Vector{Float32}}}, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}, ::AutoForwardDiff{8, Nothing}, ::Matrix{Float32})
from batched_jacobian(f::F, backend::AutoForwardDiff, x::AbstractArray) where F @ Lux /mnt/.julia/packages/Lux/L2VO7/src/autodiff/api.jl:121
Static Parameters
F = StatefulLuxLayer{Static.True, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, @NamedTuple{layer_1::@NamedTuple{weight::Matrix{Float32}, bias::Vector{Float32}}, layer_2::@NamedTuple{weight::Matrix{Float32}, bias::Vector{Float32}}}, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}
Arguments
#self#::Core.Const(Lux.batched_jacobian)
f::StatefulLuxLayer{Static.True, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, @NamedTuple{layer_1::@NamedTuple{weight::Matrix{Float32}, bias::Vector{Float32}}, layer_2::@NamedTuple{weight::Matrix{Float32}, bias::Vector{Float32}}}, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}
backend::Core.Const(AutoForwardDiff(chunksize=8))
x::Matrix{Float32}
Body::Array{Float32, 3}
1 ─ %1 = Lux.AutoDiffInternalImpl::Core.Const(Lux.AutoDiffInternalImpl)
│ %2 = Base.getproperty(%1, :batched_jacobian)::Core.Const(Lux.AutoDiffInternalImpl.batched_jacobian)
│ %3 = (%2)(f, backend, x)::Array{Float32, 3}
└── return %3
```
---
### Author:  [@liuyxpp](https://discourse.julialang.org/u/liuyxpp)
#### Post date: [April 30, 2025, 1:56am UTC](https://discourse.julialang.org/t/type-stability-of-lux-batched-jacobian/128526/5 "2025-04-30T01:56:07Z")
That’s strange. Should this be at least documented? Or better, when used in this context, `AutoForwardDiff` should have a default `chunksize`?
---
### Author:  [@NoFishLikeIan](https://discourse.julialang.org/u/NoFishLikeIan)
#### Post date: [June 10, 2025, 8:56am UTC](https://discourse.julialang.org/t/type-stability-of-lux-batched-jacobian/128526/6 "2025-06-10T08:56:01Z")
I think this just inherits the behaviour of `ForwardDiff.jl`. See [here](https://juliadiff.org/ForwardDiff.jl/v0.5/user/advanced.html#Configuring-Chunk-Size-1).