Ray loader to support mini-batch training

src/python/alias.cpp

@@ -164,6 +164,8 @@ static void set_variant(nb::args args) {
         if (curr_variant.attr("startswith")(nb::make_tuple("llvm_", "cuda_"))) {
             nb::module_ mi_python = nb::module_::import_("mitsuba.python.ad.integrators");
             nb::steal(PyImport_ReloadModule(mi_python.ptr()));
+            mi_python = nb::module_::import_("mitsuba.python.ad.loaders");
+            nb::steal(PyImport_ReloadModule(mi_python.ptr()));
         }
 
         // Only invoke user-provided callbacks after Mitsuba plugins have reloaded,

src/python/python/ad/__init__.py

@@ -1,5 +1,6 @@
 from .largesteps import LargeSteps
 from .integrators import *
+from .loaders import *
 from .optimizers import *
 from .guiding import *
 from .projective import *

src/python/python/ad/loaders/__init__.py

@@ -0,0 +1,6 @@
+import importlib
+import mitsuba as mi
+
+if mi.variant() is not None and not mi.variant().startswith('scalar'):
+    from .ray_loader import *
+    importlib.reload(ray_loader)