Making BatchLM revealed that pytorch could not handle vectorized jacfwd in contexts where there are in-place operations on arrays. This came up in the integrate_modes of bright and curvature which select pixels for adaptive upsampling. I think this can be fixed by changing the backend object fill at indices function to use the pytorch index_put function (the Jax version already works). This should be a simple fix, but then some unit tests would need to be added to stress BatchLM under various contexts (different integrate and sampling modes, as well as single model vs group model). I think it should also be tested for images that have more complex orientations than the tutorial example, which has a bunch of simple images with unit pixelscale.
It would also be good to search through all the models and check for in-place operations that happen in the brightness function (or sub parts like radial model).
Making BatchLM revealed that pytorch could not handle vectorized jacfwd in contexts where there are in-place operations on arrays. This came up in the
integrate_modes ofbrightandcurvaturewhich select pixels for adaptive upsampling. I think this can be fixed by changing the backend object fill at indices function to use the pytorch index_put function (the Jax version already works). This should be a simple fix, but then some unit tests would need to be added to stress BatchLM under various contexts (different integrate and sampling modes, as well as single model vs group model). I think it should also be tested for images that have more complex orientations than the tutorial example, which has a bunch of simple images with unit pixelscale.It would also be good to search through all the models and check for in-place operations that happen in the brightness function (or sub parts like radial model).