Refactor to ravel method

Author: Ian Schweer

pytensor/link/pytorch/dispatch/elemwise.py

@@ -1,5 +1,3 @@
-from itertools import chain
-
 import torch
 
 from pytensor.link.pytorch.dispatch.basic import pytorch_funcify
@@ -175,50 +173,32 @@ def elemwise_scalar_loop(base_fn, op, node, **kwargs):
     """
     ScalarLoop + Elemwise is too common
     to not work, but @1031, vmap won't allow it.
-    Instead, we can do the following strategy
-    1. `.unbind(dim)` will return a list of tensors
-       representing `dim` but "unwrapped". e.x.
-       ```
-       t = torch.ones(3, 4, 2)
-       len(t.unbind(0)) == 3
-       t[0].shape == torch.Size[4, 2]
-    2. If we successfully apply, the length of the list will grow
-       by the next dimension in the tensor if we flatten the previous
-       dimension result
-       ```
-       inputs = [torch.ones(3, 4, 2)]
-       level_1 = chain.from_iterable(t.unbind(0) for t in inputs)
-       level_2 = chain.from_iterable(t.unbind(0) for t in level_1)
-       len(level_2) == 3 * 4
-       ```
-    3. Eventually we'll reach single dimension tensors. At that point
-       we can iterate over each input in an element by element manner
-       and call some function
-
-    For scalar loop, we need to broadcast the tensors so all
-    the necessary values are repeated, and we "evenly" iterate through everything
+    Instead, we can ravel all the inputs, broadcasted
+    according to torch
     """
 
+    n_outputs = len(node.outputs)
+
     def elemwise_fn(*inputs):
-        Elemwise._check_runtime_broadcast(node, inputs)
-        shaped_inputs = torch.broadcast_tensors(*inputs)
-        expected_size = shaped_inputs[0].numel()
-        final_inputs = [s.clone() for s in shaped_inputs]
-        for _ in range(shaped_inputs[0].dim() - 1):
-            for i, _ in enumerate(shaped_inputs):
-                layer = chain.from_iterable([s.unbind(0) for s in final_inputs[i]])
-                final_inputs[i] = list(layer)
-
-        # make sure we still have the same number of things
-        assert len(final_inputs) == len(shaped_inputs)
-
-        # make sure each group of things are the expected size
-        assert all(len(x) == expected_size for x in final_inputs)
-
-        # make sure they are all single elements
-        assert all(len(x.shape) == 0 for tensor in final_inputs for x in tensor)
-        res = [base_fn(*args) for args in zip(*final_inputs)]
-
-        return [torch.stack(tuple(out[i] for out in res)) for i in range(len(res[0]))]
+        bcasted_inputs = torch.broadcast_tensors(*inputs)
+        raveled_inputs = [inp.ravel() for inp in bcasted_inputs]
+
+        out_shape = bcasted_inputs[0].size()
+        out_size = out_shape.numel()
+        raveled_outputs = [torch.zeros(out_size) for out in node.outputs]
+
+        for i in range(out_size):
+            core_outs = base_fn(*(inp[i] for inp in raveled_inputs))
+            if n_outputs == 1:
+                raveled_outputs[0][i] = core_outs
+            else:
+                for o in range(n_outputs):
+                    raveled_outputs[o][i] = core_outs[o]
+
+        outputs = tuple(out.view(out_shape) for out in raveled_outputs)
+        if n_outputs == 1:
+            return outputs[0]
+        else:
+            return outputs
 
     return elemwise_fn

tests/link/pytorch/test_basic.py

@@ -1,7 +1,5 @@
 from collections.abc import Callable, Iterable
 from functools import partial
-from itertools import repeat, starmap
-from unittest.mock import MagicMock, call, patch
 
 import numpy as np
 import pytest
@@ -350,6 +348,7 @@ def test_pytorch_OpFromGraph():
     f = FunctionGraph([x, y, z], [out])
     compare_pytorch_and_py(f, [xv, yv, zv])
 
+
 def test_ScalarLoop():
     n_steps = int64("n_steps")
     x0 = float64("x0")
@@ -376,25 +375,11 @@ def test_ScalarLoop_while():
     for res, expected in zip(
         [fn(n_steps=20, x0=0), fn(n_steps=20, x0=1), fn(n_steps=5, x0=1)],
         [[10, True], [10, True], [6, False]],
+        strict=True,
     ):
         np.testing.assert_allclose(res[0], np.array(expected[0]))
         np.testing.assert_allclose(res[1], np.array(expected[1]))
 
-def test_pytorch_OpFromGraph():
-    x, y, z = matrices("xyz")
-    ofg_1 = OpFromGraph([x, y], [x + y])
-    ofg_2 = OpFromGraph([x, y], [x * y, x - y])
-
-    o1, o2 = ofg_2(y, z)
-    out = ofg_1(x, o1) + o2
-
-    xv = np.ones((2, 2), dtype=config.floatX)
-    yv = np.ones((2, 2), dtype=config.floatX) * 3
-    zv = np.ones((2, 2), dtype=config.floatX) * 5
-
-    f = FunctionGraph([x, y, z], [out])
-    compare_pytorch_and_py(f, [xv, yv, zv])
-
 
 def test_ScalarLoop_Elemwise():
     n_steps = int64("n_steps")
@@ -412,26 +397,3 @@ def test_ScalarLoop_Elemwise():
     f = FunctionGraph([n_steps, x0], [state, done])
     args = [np.array(10).astype("int32"), np.arange(0, 5).astype("float32")]
     compare_pytorch_and_py(f, args)
-
-
-torch_elemwise = pytest.importorskip("pytensor.link.pytorch.dispatch.elemwise")
-
-
-@pytest.mark.parametrize("input_shapes", [[(5, 1, 1, 8), (3, 1, 1), (8,)]])
-@patch("pytensor.link.pytorch.dispatch.elemwise.Elemwise")
-def test_ScalarLoop_Elemwise_iteration_logic(_, input_shapes):
-    args = [torch.ones(*s) for s in input_shapes[:-1]] + [
-        torch.zeros(*input_shapes[-1])
-    ]
-    mock_inner_func = MagicMock()
-    ret_value = torch.rand(2, 2).unbind(0)
-    mock_inner_func.f.return_value = ret_value
-    elemwise_fn = torch_elemwise.elemwise_scalar_loop(mock_inner_func.f, None, None)
-    result = elemwise_fn(*args)
-    for actual, expected in zip(ret_value, result):
-        assert torch.all(torch.eq(*torch.broadcast_tensors(actual, expected)))
-    np.testing.assert_equal(mock_inner_func.f.call_count, len(result[0]))
-
-    expected_args = torch.FloatTensor([1.0] * (len(input_shapes) - 1) + [0.0]).unbind(0)
-    expected_calls = starmap(call, repeat(expected_args, mock_inner_func.f.call_count))
-    mock_inner_func.f.assert_has_calls(expected_calls)