Support for unfold (and fold) with Lux.jl & Reactant.jl

[bvdmitri]

Lux documentation mentions NNlib.unfold operation, but the current implementation within NNlib uses @threads thus is not compatible with @compile macro from Reactant.jl

[avik-pal]

jax-ml/jax#5968

[bvdmitri]

Nice! Looking forward for the PR to be merged!

[avik-pal]

Will cut a release once #1737 is merged

[bartvanerp]

Hi @avik-pal , just wanted to check the latests PR and swapped our custom unfolding through convolution (as in the upvoted answer in jax-ml/jax#5968 ) by NNlib.unfold(), but it started throwing the error nested task error: No MLIR context is active when calling unfold with Reactant. Any idea what might be causing this?

[avik-pal]

can you share a runnable example?

The error essentially is saying we are trying to do operations on TracedRArrays outside of @compile. But that should not happen.

[bartvanerp]

Sorry for the noise, seemed to be caused by a local installation of NNlib (and lack of sleep :) )