Update relevant imports to pytensor.graph.traversal

Author: jessegrabowski

pymc/sampling/jax.py

@@ -31,9 +31,9 @@
 from jax.lax import scan
 from numpy.typing import ArrayLike
 from pytensor.compile import SharedVariable, mode
-from pytensor.graph.basic import graph_inputs
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.replace import clone_replace
+from pytensor.graph.traversal import graph_inputs
 from pytensor.link.jax.dispatch import jax_funcify
 from pytensor.raise_op import Assert
 from pytensor.tensor import TensorVariable

pymc/testing.py

@@ -26,8 +26,9 @@
 from numpy import testing as npt
 from pytensor.compile import SharedVariable
 from pytensor.compile.mode import Mode
-from pytensor.graph.basic import Constant, Variable, equal_computations, graph_inputs
+from pytensor.graph.basic import Constant, Variable, equal_computations
 from pytensor.graph.rewriting.basic import in2out
+from pytensor.graph.traversal import graph_inputs
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.random.type import RandomType

tests/logprob/test_basic.py

@@ -42,7 +42,8 @@
 import pytest
 import scipy.stats.distributions as sp
 
-from pytensor.graph.basic import ancestors, equal_computations
+from pytensor.graph.basic import equal_computations
+from pytensor.graph.traversal import ancestors
 from pytensor.tensor.random.op import RandomVariable
 from scipy import stats
 

tests/logprob/test_utils.py

@@ -41,7 +41,8 @@
 from pytensor import function
 from pytensor import tensor as pt
 from pytensor.compile import get_default_mode
-from pytensor.graph.basic import ancestors, equal_computations
+from pytensor.graph.basic import equal_computations
+from pytensor.graph.traversal import ancestors
 from pytensor.tensor.random.basic import NormalRV
 from pytensor.tensor.random.op import RandomVariable
 

tests/model/test_core.py

@@ -33,7 +33,7 @@
 import scipy.stats as st
 
 from pytensor.graph import graph_inputs
-from pytensor.graph.basic import get_var_by_name
+from pytensor.graph.traversal import get_var_by_name
 from pytensor.raise_op import Assert
 from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.variable import TensorConstant

tests/sampling/test_forward.py

@@ -28,7 +28,7 @@
 from pytensor import Mode, shared
 from pytensor.compile import SharedVariable
 from pytensor.graph import graph_inputs
-from pytensor.graph.basic import get_var_by_name, variable_depends_on
+from pytensor.graph.traversal import get_var_by_name, variable_depends_on
 from pytensor.tensor.variable import TensorConstant
 from scipy import stats