Preserve numpy < 2.0 Unique inverse output shape

In numpy 2.0, if axis=None, then np.unique
does not flatten the inverse indices returned
if return_inverse=True

A helper function has been added to npy_2_compat.py
to mimic the output of `np.unique` from version of
numpy before 2.0

Author: brendan-m-murphy

pytensor/npy_2_compat.py

@@ -63,6 +63,28 @@
 numpy_maxdims = 64 if using_numpy_2 else 32
 
 
+# function that replicates np.unique from numpy < 2.0
+def old_np_unique(
+    arr, return_index=False, return_inverse=False, return_counts=False, axis=None
+):
+    """Replicate np.unique from numpy versions < 2.0"""
+    if not return_inverse or not using_numpy_2:
+        return np.unique(arr, return_index, return_inverse, return_counts, axis)
+
+    outs = list(np.unique(arr, return_index, return_inverse, return_counts, axis))
+
+    inv_idx = 2 if return_index else 1
+
+    if axis is None:
+        outs[inv_idx] = np.ravel(outs[inv_idx])
+    else:
+        inv_shape = (arr.shape[axis],)
+        outs[inv_idx] = outs[inv_idx].reshape(inv_shape)
+
+    return tuple(outs)
+
+
+# compatibility header for C code
 def npy_2_compat_header() -> str:
     """Compatibility header that Numpy suggests is vendored with code that uses Numpy < 2.0 and Numpy 2.x"""
     return dedent("""

pytensor/tensor/extra_ops.py

@@ -20,6 +20,7 @@
     normalize_axis_index,
     npy_2_compat_header,
     numpy_axis_is_none_flag,
+    old_np_unique,
 )
 from pytensor.raise_op import Assert
 from pytensor.scalar import int64 as int_t
@@ -1226,6 +1227,9 @@ class Unique(Op):
     """
     Wraps `numpy.unique`.
 
+    The indices returned when `return_inverse` is True are ravelled
+    to match the behavior of `numpy.unique` from before numpy version 2.0.
+
     Examples
     --------
     >>> import numpy as np
@@ -1271,17 +1275,21 @@ def make_node(self, x):
 
         outputs = [TensorType(dtype=x.dtype, shape=out_shape)()]
         typ = TensorType(dtype="int64", shape=(None,))
+
         if self.return_index:
             outputs.append(typ())
+
         if self.return_inverse:
             outputs.append(typ())
+
         if self.return_counts:
             outputs.append(typ())
+
         return Apply(self, [x], outputs)
 
     def perform(self, node, inputs, output_storage):
         [x] = inputs
-        outs = np.unique(
+        outs = old_np_unique(
             x,
             return_index=self.return_index,
             return_inverse=self.return_inverse,
@@ -1306,9 +1314,14 @@ def infer_shape(self, fgraph, node, i0_shapes):
             out_shapes[0] = tuple(shape)
 
         if self.return_inverse:
-            shape = prod(x_shape) if self.axis is None else x_shape[axis]
             return_index_out_idx = 2 if self.return_index else 1
-            out_shapes[return_index_out_idx] = (shape,)
+
+            if self.axis is not None:
+                shape = (x_shape[axis],)
+            else:
+                shape = (prod(x_shape),)
+
+            out_shapes[return_index_out_idx] = shape
 
         return out_shapes
 

tests/tensor/test_extra_ops.py

@@ -9,6 +9,7 @@
 from pytensor.compile.mode import Mode
 from pytensor.configdefaults import config
 from pytensor.graph.basic import Constant, applys_between, equal_computations
+from pytensor.npy_2_compat import old_np_unique
 from pytensor.raise_op import Assert
 from pytensor.tensor import alloc
 from pytensor.tensor.elemwise import DimShuffle
@@ -899,18 +900,19 @@ def setup_method(self):
     )
     def test_basic_vector(self, x, inp, axis):
         list_outs_expected = [
-            np.unique(inp, axis=axis),
-            np.unique(inp, True, axis=axis),
-            np.unique(inp, False, True, axis=axis),
-            np.unique(inp, True, True, axis=axis),
-            np.unique(inp, False, False, True, axis=axis),
-            np.unique(inp, True, False, True, axis=axis),
-            np.unique(inp, False, True, True, axis=axis),
-            np.unique(inp, True, True, True, axis=axis),
+            old_np_unique(inp, axis=axis),
+            old_np_unique(inp, True, axis=axis),
+            old_np_unique(inp, False, True, axis=axis),
+            old_np_unique(inp, True, True, axis=axis),
+            old_np_unique(inp, False, False, True, axis=axis),
+            old_np_unique(inp, True, False, True, axis=axis),
+            old_np_unique(inp, False, True, True, axis=axis),
+            old_np_unique(inp, True, True, True, axis=axis),
         ]
         for params, outs_expected in zip(
             self.op_params, list_outs_expected, strict=True
         ):
+            print(params)
             out = pt.unique(x, *params, axis=axis)
             f = pytensor.function(inputs=[x], outputs=out)
             outs = f(inp)