BlockCD does not support non-jitted functions

google / jaxopt

Hardware accelerated, batchable and differentiable optimizers in JAX.

https://jaxopt.github.io

Apache License 2.0

936 stars 66 forks source link

BlockCD does not support non-jitted functions #462

Open zaccharieramzi opened 1 year ago

zaccharieramzi commented 1 year ago

For now because fori_loop automatically compiles the function, BlockCD does not support non-jitted functions. The resolution should be very similar to what was done in #449 .