[fault1]

yes, see: https://juliadiff.org/ChainRulesCore.jl/dev/maths/nondiff_po...

and related neat usages in set based optimization methods in MathOptInterface (part of JuMP.jl): https://matbesancon.xyz/post/2020-12-24-chains_sets2/

[ogogmad]

To me, it looks like you're computing a single element from the set, instead of the whole set. Which makes sense, but it's an important caveat.