run pre-commit hooks

Author: jessegrabowski

pytensor/link/jax/dispatch/scan.py

@@ -3,7 +3,7 @@
 
 from pytensor.link.jax.dispatch.basic import jax_funcify
 from pytensor.scan.op import Scan
-from pytensor.compile.sharedvalue import SharedVariable
+
 
 @jax_funcify.register(Scan)
 def jax_funcify_Scan(op: Scan, **kwargs):
@@ -42,11 +42,12 @@ def scan(*outer_inputs):
             op.outer_non_seqs(outer_inputs),
         )  # JAX `init`
 
-        ndim_out_core = [getattr(x, 'ndim', 0) for x in op.fgraph.outputs]
-        ndim_in_core  = [getattr(x, 'ndim', 0) for x in mit_sot_init + sit_sot_init]
+        ndim_out_core = [getattr(x, "ndim", 0) for x in op.fgraph.outputs]
+        ndim_in_core = [getattr(x, "ndim", 0) for x in mit_sot_init + sit_sot_init]
 
-        n_batchdims = [int(dims_out > 0) for dims_out, dims_in in zip(ndim_out_core, ndim_in_core)] +\
-            [0] * op.info.n_nit_sot
+        n_batchdims = [
+            int(dims_out > 0) for dims_out, dims_in in zip(ndim_out_core, ndim_in_core)
+        ] + [0] * op.info.n_nit_sot
 
         def jax_args_to_inner_func_args(carry, x):
             """Convert JAX scan arguments into format expected by scan_inner_func.
@@ -157,13 +158,17 @@ def get_partial_traces(traces):
                 + op.outer_nitsot(outer_inputs)
             )
             partial_traces = []
-            for init_state, trace, n_batchdim, buffer in zip(init_states, traces, n_batchdims, buffers):
+            for init_state, trace, n_batchdim, buffer in zip(
+                init_states, traces, n_batchdims, buffers
+            ):
                 if init_state is not None:
                     # MIT-SOT and SIT-SOT: The final output should be as long as the input buffer
                     batch_idx = range(n_batchdim)
                     full_trace = jnp.concatenate(
-                        [jnp.atleast_1d(jnp.expand_dims(init_state, batch_idx)), 
-                        jnp.atleast_1d(trace)],
+                        [
+                            jnp.atleast_1d(jnp.expand_dims(init_state, batch_idx)),
+                            jnp.atleast_1d(trace),
+                        ],
                         axis=0,
                     )
                     buffer_size = buffer.shape[0]

tests/link/jax/test_scan.py

@@ -13,14 +13,11 @@
 from pytensor.scan.op import Scan
 from pytensor.tensor import random
 from pytensor.tensor.math import gammaln, log
-from pytensor.tensor.type import lscalar, scalar, vector, dmatrix, dvector
+from pytensor.tensor.type import dmatrix, dvector, lscalar, scalar, vector
 from tests.link.jax.test_basic import compare_jax_and_py
 
 
-# jax = pytest.importorskip("jax")
-
-import jax
-jax.config.update('jax_platform_name', 'cpu')
+jax = pytest.importorskip("jax")
 
 
 @pytest.mark.parametrize("view", [None, (-1,), slice(-2, None, None)])
@@ -322,22 +319,24 @@ def input_step_fn(y_tm1, y_tm3, a):
     compare_jax_and_py(out_fg, test_input_vals)
 
 
-@pytest.mark.parametrize('x0_func', [dvector, dmatrix])
-@pytest.mark.parametrize('A_func',  [dmatrix, dmatrix])
+@pytest.mark.parametrize("x0_func", [dvector, dmatrix])
+@pytest.mark.parametrize("A_func", [dmatrix, dmatrix])
 def test_nd_scan_sit_sot(x0_func, A_func):
-    x0 = x0_func('x0')
-    A  = A_func('A')
+    x0 = x0_func("x0")
+    A = A_func("A")
 
     n_steps = 3
     k = 3
 
     # Must specify mode = JAX for the inner func to avoid a GEMM Op in the JAX graph
-    xs, _ = scan(lambda X, A: A @ X, 
-                            non_sequences=[A],
-                            outputs_info=[x0],
-                            n_steps=n_steps,
-                            mode=get_mode('JAX'))
-                
+    xs, _ = scan(
+        lambda X, A: A @ X,
+        non_sequences=[A],
+        outputs_info=[x0],
+        n_steps=n_steps,
+        mode=get_mode("JAX"),
+    )
+
     x0_val = np.arange(k) if x0.ndim == 1 else np.diag(np.arange(k))
     A_val = np.eye(k)
 
@@ -347,19 +346,21 @@ def test_nd_scan_sit_sot(x0_func, A_func):
 
 
 def test_nd_scan_sit_sot_with_seq():
-    x = dmatrix('x0')
-    A  = dmatrix('A')
+    x = dmatrix("x0")
+    A = dmatrix("A")
 
     n_steps = 3
     k = 3
 
     # Must specify mode = JAX for the inner func to avoid a GEMM Op in the JAX graph
-    xs, _ = scan(lambda X, A: A @ X, 
-                            non_sequences=[A],
-                            sequences=[x],
-                            n_steps=n_steps,
-                            mode=get_mode('JAX'))
-                
+    xs, _ = scan(
+        lambda X, A: A @ X,
+        non_sequences=[A],
+        sequences=[x],
+        n_steps=n_steps,
+        mode=get_mode("JAX"),
+    )
+
     x_val = np.tile(np.arange(k), n_steps).reshape(n_steps, k)
     A_val = np.eye(k)
 
@@ -369,17 +370,17 @@ def test_nd_scan_sit_sot_with_seq():
 
 
 def test_nd_scan_mit_sot():
-    x0 = dmatrix('x0')
-    A  = dmatrix('A')
-    B  = dmatrix('B')
+    x0 = dmatrix("x0")
+    A = dmatrix("A")
+    B = dmatrix("B")
 
     # Must specify mode = JAX for the inner func to avoid a GEMM Op in the JAX graph
     xs, _ = scan(
         lambda xtm3, xtm1, A, B: A @ xtm3 + B @ xtm1,
         outputs_info=[{"initial": x0, "taps": [-3, -1]}],
         non_sequences=[A, B],
         n_steps=10,
-        mode = get_mode('JAX')
+        mode=get_mode("JAX"),
     )
 
     fg = FunctionGraph([x0, A, B], [xs])
@@ -392,8 +393,8 @@ def test_nd_scan_mit_sot():
 
 
 def test_nd_scan_sit_sot_with_carry():
-    x0 = dvector('x0')
-    A  = dmatrix('A')
+    x0 = dvector("x0")
+    A = dmatrix("A")
 
     def step(x, A):
         return A @ x, x.sum()
@@ -404,7 +405,7 @@ def step(x, A):
         outputs_info=[x0, None],
         non_sequences=[A],
         n_steps=10,
-        mode = get_mode('JAX')
+        mode=get_mode("JAX"),
     )
 
     fg = FunctionGraph([x0, A], xs)