Fixed JAX unit tests for new structure

Author: HangenYuu

pytensor/link/jax/dispatch/__init__.py

@@ -2,18 +2,20 @@
 from pytensor.link.jax.dispatch.basic import jax_funcify, jax_typify
 
 # Load dispatch specializations
-import pytensor.link.jax.dispatch.scalar
-import pytensor.link.jax.dispatch.tensor_basic
-import pytensor.link.jax.dispatch.subtensor
-import pytensor.link.jax.dispatch.shape
+import pytensor.link.jax.dispatch.blas
+import pytensor.link.jax.dispatch.blockwise
+import pytensor.link.jax.dispatch.elemwise
 import pytensor.link.jax.dispatch.extra_ops
+import pytensor.link.jax.dispatch.math
 import pytensor.link.jax.dispatch.nlinalg
-import pytensor.link.jax.dispatch.slinalg
 import pytensor.link.jax.dispatch.random
-import pytensor.link.jax.dispatch.elemwise
+import pytensor.link.jax.dispatch.scalar
 import pytensor.link.jax.dispatch.scan
-import pytensor.link.jax.dispatch.sparse
-import pytensor.link.jax.dispatch.blockwise
+import pytensor.link.jax.dispatch.shape
+import pytensor.link.jax.dispatch.slinalg
 import pytensor.link.jax.dispatch.sort
+import pytensor.link.jax.dispatch.sparse
+import pytensor.link.jax.dispatch.subtensor
+import pytensor.link.jax.dispatch.tensor_basic
 
 # isort: on

pytensor/link/jax/dispatch/math.py

@@ -1,4 +1,5 @@
 import jax.numpy as jnp
+import numpy as np
 
 from pytensor.link.jax.dispatch import jax_funcify
 from pytensor.tensor.math import Argmax, Dot, Max
@@ -36,12 +37,10 @@ def argmax(x):
 
         # NumPy does not support multiple axes for argmax; this is a
         # work-around
-        keep_axes = jnp.array(
-            [i for i in range(x.ndim) if i not in axes], dtype="int64"
-        )
+        keep_axes = np.array([i for i in range(x.ndim) if i not in axes], dtype="int64")
         # Not-reduced axes in front
         transposed_x = jnp.transpose(
-            x, jnp.concatenate((keep_axes, jnp.array(axes, dtype="int64")))
+            x, tuple(np.concatenate((keep_axes, np.array(axes, dtype="int64"))))
         )
         kept_shape = transposed_x.shape[: len(keep_axes)]
         reduced_shape = transposed_x.shape[len(keep_axes) :]
@@ -50,9 +49,9 @@ def argmax(x):
         # Otherwise reshape would complain citing float arg
         new_shape = (
             *kept_shape,
-            jnp.prod(jnp.array(reduced_shape, dtype="int64"), dtype="int64"),
+            np.prod(np.array(reduced_shape, dtype="int64"), dtype="int64"),
         )
-        reshaped_x = transposed_x.reshape(new_shape)
+        reshaped_x = transposed_x.reshape(tuple(new_shape))
 
         max_idx_res = jnp.argmax(reshaped_x, axis=-1).astype("int64")