Skip to content

Fix #1742

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix #1742

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 0 additions & 22 deletions pytensor/npy_2_compat.py
View file
Open in desktop

This file was deleted.

35 changes: 20 additions & 15 deletions pytensor/tensor/extra_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
from pytensor.link.c.op import COp
from pytensor.link.c.params_type import ParamsType
from pytensor.link.c.type import EnumList, Generic
from pytensor.npy_2_compat import old_np_unique
from pytensor.raise_op import Assert
from pytensor.scalar import int64 as int_t
from pytensor.scalar import upcast
Expand All @@ -37,7 +36,6 @@
lt,
maximum,
minimum,
prod,
sign,
switch,
)
Expand Down Expand Up @@ -1204,31 +1202,36 @@ def make_node(self, x):
if axis is None:
out_shape = (None,)
else:
if axis >= x.type.ndim:
raise ValueError(
f"Axis {axis} out of range for input {x} with ndim={x.type.ndim}."
)
out_shape = tuple(
None if dim == axis else s for dim, s in enumerate(x.type.shape)
)

outputs = [TensorType(dtype=x.dtype, shape=out_shape)()]
typ = TensorType(dtype="int64", shape=(None,))

index_type = TensorType("int64", shape=(None,))
count_type = TensorType("int64", shape=(None,))

# inverse gets NumPy-2 behavior
if axis is None:
inverse_shape = x.type.shape # same as input
else:
inverse_shape = (x.type.shape[axis],)

inverse_type = TensorType("int64", shape=inverse_shape)

if self.return_index:
outputs.append(typ())
outputs.append(index_type())

if self.return_inverse:
outputs.append(typ())
outputs.append(inverse_type())

if self.return_counts:
outputs.append(typ())
outputs.append(count_type())

return Apply(self, [x], outputs)

def perform(self, node, inputs, output_storage):
[x] = inputs
outs = old_np_unique(
outs = np.unique(
x,
return_index=self.return_index,
return_inverse=self.return_inverse,
Expand All @@ -1255,10 +1258,12 @@ def infer_shape(self, fgraph, node, i0_shapes):
if self.return_inverse:
return_index_out_idx = 2 if self.return_index else 1

if self.axis is not None:
shape = (x_shape[axis],)
if self.axis is None:
# NumPy 2: inverse has same shape as input
shape = x_shape
else:
shape = (prod(x_shape),)
# NumPy 2: inverse has same shape as reduction axis
shape = (x_shape[self.axis],)

out_shapes[return_index_out_idx] = shape

Expand Down
17 changes: 8 additions & 9 deletions tests/tensor/test_extra_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
from pytensor.graph import rewrite_graph
from pytensor.graph.basic import Constant, equal_computations
from pytensor.graph.traversal import applys_between
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
Expand Down Expand Up @@ -902,14 +901,14 @@ def setup_method(self):
)
def test_basic_vector(self, x, inp, axis):
list_outs_expected = [
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),
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),
]
for params, outs_expected in zip(
self.op_params, list_outs_expected, strict=True
Expand Down
Loading