From eab8addcd9eaf212dd828e14f81d610f7b5c4d07 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 14 Nov 2024 11:00:28 -0500 Subject: [PATCH] fix: don't preserve structure in checks --- src/contrib/contrib.jl | 2 +- src/contrib/debug.jl | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/contrib/contrib.jl b/src/contrib/contrib.jl index 3e62563f32..a4b07170b8 100644 --- a/src/contrib/contrib.jl +++ b/src/contrib/contrib.jl @@ -6,7 +6,7 @@ using ChainRulesCore: ChainRulesCore using Compat: @compat using ConcreteStructs: @concrete using FastClosures: @closure -using Functors: Functors, KeyPath, fmap_with_path, functor +using Functors: Functors, KeyPath, fmap_with_path, fmapstructure_with_path, functor using Markdown: @doc_str using Optimisers: Optimisers using Random: AbstractRNG, Random diff --git a/src/contrib/debug.jl b/src/contrib/debug.jl index 7d7b388618..38c13f36fe 100644 --- a/src/contrib/debug.jl +++ b/src/contrib/debug.jl @@ -91,7 +91,8 @@ function check_nan_and_throw(x, str::AbstractString, layer, location::KeyPath) return x end - return fmap_with_path(nan_check, x) + fmapstructure_with_path(nan_check, x) + return end function debug_layer_impl(layer, x, ps, st, location, error_check, _)