Implement _sum method fo jax

NVIDIA-Merlin · Feb 20, 2023 · 1242ecd · 1242ecd
1 parent 3fe94b6
commit 1242ecd
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/merlin/dataloader/jax.py b/merlin/dataloader/jax.py
@@ -69,6 +69,9 @@ def _split_fn(self, tensor, idx, axis=0):
 
     _tensor_split = _split_fn
 
+    def _sum(self, tensor):
+        return tensor.sum()
+
     def _to_tensor(self, gdf):
         if gdf.empty:
             return