[ML][docs] Add documentation regarding mapping in RDataLoader.

[JazzyJonah]

This Pull request adds a section in the RDataLoader documentation detailing how to transform your data after an RDataLoader export.

• updated the docs (if necessary)

[github-actions]

Test Results

    23 files      23 suites   3d 15h 53m 27s ⏱️
 3 882 tests  3 881 ✅ 0 💤 1 ❌
79 059 runs  79 057 ✅ 1 💤 1 ❌

For more details on these failures, see this check.

Results for commit d7e5394.

♻️ This comment has been updated with latest results.

[guitargeek]

Hi @JazzyJonah, thanks a lot for this, and for the really thorough tests! 🙏

After playing with it a bit, I actually think we might not need the transform= parameter at all: the exact same thing is already achievable with the builtin map, e.g.

for x, y in map(reshape, dl.as_numpy()):
    ...

This yields identical batches (I checked for numpy, torch, and jax). The only subtlety is that a map(...) object is single-use, so it goes inside the epoch loop rather than being stored once. But that's the natural place for it anyway, and a short note in the docstrings would cover it nicely.

The nice thing is that this means the capability is essentially already there for free. I'd lean towards not adding the extra argument: every new parameter is a bit of maintenance liability, and it makes the as_numpy/as_torch/as_jax signatures and docs a little less succinct. Relying on map keeps the API surface small while still giving users the full flexibility.

So for this PR, that would mean trading the implementation for a small documentation addition, and dropping the transform tests. Chaining array/tensor generators and functions with map is a well-proven idiom and already outside the responsibility of ROOT, so I'm afraid testing the map approach would be unnecessary churn.

Here is how the tests would look like using map:

diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader.py b/bindings/pyroot/pythonizations/test/ml_dataloader.py
index 6427906cf0c..1b4782e6911 100644
--- a/bindings/pyroot/pythonizations/test/ml_dataloader.py
+++ b/bindings/pyroot/pythonizations/test/ml_dataloader.py
@@ -1226,11 +1226,9 @@ class DataLoaderMultipleFiles(unittest.TestCase):
             def reshape_bad(batch, newShape):
                 return np.reshape(batch, newShape)
 
-            arrs_reshape = dl.as_numpy(transform=reshape)
-            arrs_reshape_default = dl.as_numpy(transform=reshape_default)
-            with self.assertRaises(TypeError):
-                dl.as_numpy(transform=reshape_bad)
-            for arrs in (arrs_reshape, arrs_reshape_default):
+            arrs_reshape = dl.as_numpy()
+            arrs_reshape_default = dl.as_numpy()
+            for arrs in (map(reshape, arrs_reshape), map(reshape_default, arrs_reshape_default)):
                 for arr in arrs:
                     self.assertEqual(len(arr), 2)
                     x, y = arr
@@ -1274,11 +1272,9 @@ class DataLoaderMultipleFiles(unittest.TestCase):
             def reshape_bad(batch, newShape):
                 return tuple(torch.reshape(tensor, newShape) for tensor in batch)
 
-            tensors_reshape = dl.as_torch(transform=reshape)
-            tensors_reshape_default = dl.as_torch(transform=reshape_default)
-            with self.assertRaises(TypeError):
-                dl.as_torch(transform=reshape_bad)
-            for tensors in (tensors_reshape, tensors_reshape_default):
+            tensors_reshape = dl.as_torch()
+            tensors_reshape_default = dl.as_torch()
+            for tensors in (map(reshape, tensors_reshape), map(reshape_default, tensors_reshape_default)):
                 for tensor in tensors:
                     self.assertEqual(len(tensor), 2)
                     x, y = tensor
@@ -1322,11 +1318,9 @@ class DataLoaderMultipleFiles(unittest.TestCase):
             def reshape_bad(batch, newShape):
                 return tuple(jnp.reshape(arr, newShape) for arr in batch)
 
-            arrs_reshape = dl.as_jax(transform=reshape)
-            arrs_reshape_default = dl.as_jax(transform=reshape_default)
-            with self.assertRaises(TypeError):
-                dl.as_jax(transform=reshape_bad)
-            for arrs in (arrs_reshape, arrs_reshape_default):
+            arrs_reshape = dl.as_jax()
+            arrs_reshape_default = dl.as_jax()
+            for arrs in (map(reshape, arrs_reshape), map(reshape_default, arrs_reshape_default)):
                 for arr in arrs:
                     self.assertEqual(len(arr), 2)
                     x, y = arr
@@ -2461,11 +2455,9 @@ class DataLoaderEagerLoading(unittest.TestCase):
             def reshape_bad(batch, newShape):
                 return np.reshape(batch, newShape)
 
-            arrs_reshape = dl.as_numpy(transform=reshape)
-            arrs_reshape_default = dl.as_numpy(transform=reshape_default)
-            with self.assertRaises(TypeError):
-                dl.as_numpy(transform=reshape_bad)
-            for arrs in (arrs_reshape, arrs_reshape_default):
+            arrs_reshape = dl.as_numpy()
+            arrs_reshape_default = dl.as_numpy()
+            for arrs in (map(reshape, arrs_reshape), map(reshape_default, arrs_reshape_default)):
                 for arr in arrs:
                     self.assertEqual(len(arr), 2)
                     x, y = arr
@@ -2505,11 +2497,9 @@ class DataLoaderEagerLoading(unittest.TestCase):
             def reshape_bad(batch, newShape):
                 return tuple(torch.reshape(tensor, newShape) for tensor in batch)
 
-            tensors_reshape = dl.as_torch(transform=reshape)
-            tensors_reshape_default = dl.as_torch(transform=reshape_default)
-            with self.assertRaises(TypeError):
-                dl.as_torch(transform=reshape_bad)
-            for tensors in (tensors_reshape, tensors_reshape_default):
+            tensors_reshape = dl.as_torch()
+            tensors_reshape_default = dl.as_torch()
+            for tensors in (map(reshape, tensors_reshape), map(reshape_default, tensors_reshape_default)):
                 for tensor in tensors:
                     self.assertEqual(len(tensor), 2)
                     x, y = tensor
@@ -2549,11 +2539,9 @@ class DataLoaderEagerLoading(unittest.TestCase):
             def reshape_bad(batch, newShape):
                 return tuple(jnp.reshape(arr, newShape) for arr in batch)
 
-            arrs_reshape = dl.as_jax(transform=reshape)
-            arrs_reshape_default = dl.as_jax(transform=reshape_default)
-            with self.assertRaises(TypeError):
-                dl.as_jax(transform=reshape_bad)
-            for arrs in (arrs_reshape, arrs_reshape_default):
+            arrs_reshape = dl.as_jax()
+            arrs_reshape_default = dl.as_jax()
+            for arrs in (map(reshape, arrs_reshape), map(reshape_default, arrs_reshape_default)):
                 for arr in arrs:
                     self.assertEqual(len(arr), 2)
                     x, y = arr

[guitargeek]

Thank you so much, this is a very clean solution now 👍 👍

[siliataider]

Very nice @JazzyJonah thanks!

[vepadulano]

Nice!