Currently several places in the codebase uses this pattern:

adapt(eltypeθ,
   hcat(vec(map(points -> collect(points), Iterators.product(span...)))...))

But, adapt doesn't work with eltype, for example:

julia> adapt(Float32, rand(2, 3))
2×3 Matrix{Float64}:
 0.0893353  0.197801   0.457015
 0.796801   0.0334906  0.255498

This causes adapt to simply ignore the eltype and constructs non-concrete array types, which eventually leads to issues like https://github.com/SciML/NeuralPDE.jl/actions/runs/10485955321/job/29043160310?pr=882#step:6:1231 (matrix multiplication with non-concrete elements).

Lux does have mechanisms to catch these. While not perfect, activating the following preferences will throw errors whenever a "bad" type is being propagated. I don't turn these on by default, since they can have some overhead but for testing purposes these should be set to "error" on CI.

using Preferences

Preferences.set_preferences!("Lux", "eltype_mismatch_handling" => "error")

Preferences.set_preferences!("LuxLib", "instability_check" => "error")
Preferences.set_preferences!("LuxCore", "instability_check" => "error")