Revert "Use non recursive algorithm in rebuild_collect_shared"

This reverts commit ae43c14.

Breaks the @pytensor_jit example

Author: ricardoV94

pytensor/compile/function/pfunc.py

@@ -179,54 +179,47 @@ def clone_v_get_shared_updates(v, copy_inputs_over):
 
         """
         # this co-recurses with clone_a
-        stack = [v]
-        try:
-            while True:
-                v = stack.pop()
-                if v in clone_d:
-                    continue
-                if (apply := v.owner) is not None:
-                    if all(i in clone_d for i in apply.inputs):
-                        # all inputs have been cloned, we can clone this node
-                        clone_node_and_cache(
-                            apply,
-                            clone_d,
-                            strict=rebuild_strict,
-                            clone_inner_graphs=clone_inner_graphs,
+        assert v is not None
+        if v in clone_d:
+            return clone_d[v]
+        if v.owner:
+            owner = v.owner
+            if owner not in clone_d:
+                for i in owner.inputs:
+                    clone_v_get_shared_updates(i, copy_inputs_over)
+                clone_node_and_cache(
+                    owner,
+                    clone_d,
+                    strict=rebuild_strict,
+                    clone_inner_graphs=clone_inner_graphs,
+                )
+            return clone_d.setdefault(v, v)
+        elif isinstance(v, SharedVariable):
+            if v not in shared_inputs:
+                shared_inputs.append(v)
+            if v.default_update is not None:
+                # Check that v should not be excluded from the default
+                # updates list
+                if no_default_updates is False or (
+                    isinstance(no_default_updates, list) and v not in no_default_updates
+                ):
+                    # Do not use default_update if a "real" update was
+                    # provided
+                    if v not in update_d:
+                        v_update = v.type.filter_variable(
+                            v.default_update, allow_convert=False
                         )
-                    else:
-                        # expand on the inputs
-                        stack.extend(apply.inputs)
-                else:
-                    clone_d[v] = v if copy_inputs_over else v.clone()
-
-                    # Special handling of SharedVariables
-                    if isinstance(v, SharedVariable):
-                        if v not in shared_inputs:
-                            shared_inputs.append(v)
-                        if v.default_update is not None:
-                            # Check that v should not be excluded from the default
-                            # updates list
-                            if no_default_updates is False or (
-                                isinstance(no_default_updates, list)
-                                and v not in no_default_updates
-                            ):
-                                # Do not use default_update if a "real" update was
-                                # provided
-                                if v not in update_d:
-                                    v_update = v.type.filter_variable(
-                                        v.default_update, allow_convert=False
-                                    )
-                                    if not v.type.is_super(v_update.type):
-                                        raise TypeError(
-                                            "An update must have a type compatible with "
-                                            "the original shared variable"
-                                        )
-                                    update_d[v] = v_update
-                                    update_expr.append((v, v_update))
-        except IndexError:
-            pass  # stack is empty
-        return clone_d[v]
+                        if not v.type.is_super(v_update.type):
+                            raise TypeError(
+                                "An update must have a type compatible with "
+                                "the original shared variable"
+                            )
+                        update_d[v] = v_update
+                        update_expr.append((v, v_update))
+        if not copy_inputs_over:
+            return clone_d.setdefault(v, v.clone())
+        else:
+            return clone_d.setdefault(v, v)
 
     # initialize the clone_d mapping with the replace dictionary
     if replace is None: