Updated 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

The output shape has been changed to match numpy,
dependent on the version of numpy.

Author: brendan-m-murphy

pytensor/tensor/extra_ops.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 from numpy.exceptions import AxisError
+from numpy.lib import NumpyVersion
 
 
 try:
@@ -1241,12 +1242,25 @@ 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 NumpyVersion(np.__version__) >= "2.0.0rc1":
+                if axis is None:
+                    inverse_shape = TensorType(dtype="int64", shape=x.type.shape)
+                else:
+                    inverse_shape = TensorType(
+                        dtype="int64", shape=(x.type.shape[axis],)
+                    )
+                outputs.append(inverse_shape())
+            else:
+                outputs.append(typ())
+
         if self.return_counts:
             outputs.append(typ())
+
         return Apply(self, [x], outputs)
 
     def perform(self, node, inputs, output_storage):
@@ -1276,9 +1290,16 @@ 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],)
+            elif NumpyVersion(np.__version__) >= "2.0.0rc1":
+                shape = x_shape
+            else:
+                shape = (prod(x_shape),)
+
+            out_shapes[return_index_out_idx] = shape
 
         return out_shapes