Skip to content

Rewrite reshapes that only expand or squeeze dims #1200

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.

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

Merged
merged 6 commits into from
Feb 17, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
10 changes: 0 additions & 10 deletions pytensor/graph/rewriting/basic.py
Original file line number Diff line number Diff line change
Expand Up @@ -2800,16 +2800,6 @@ def _check_chain(r, chain):
return r is not None


def check_chain(r, *chain):
"""
WRITEME

"""
if isinstance(r, Apply):
r = r.outputs[0]
return _check_chain(r, reduce(list.__iadd__, ([x, 0] for x in chain)))


def pre_greedy_node_rewriter(
fgraph: FunctionGraph, rewrites: Sequence[NodeRewriter], out: Variable
) -> Variable:
Expand Down
26 changes: 15 additions & 11 deletions pytensor/tensor/elemwise.py
Original file line number Diff line number Diff line change
Expand Up @@ -166,15 +166,20 @@ def __init__(self, *, input_ndim: int, new_order: Sequence[int | Literal["x"]]):
self.transposition = self.shuffle + drop
# List of dimensions of the output that are broadcastable and were not
# in the original input
self.augment = sorted(i for i, x in enumerate(new_order) if x == "x")
self.augment = augment = sorted(i for i, x in enumerate(new_order) if x == "x")
self.drop = drop

self.is_left_expand_dims = self.augment and (
dims_are_shuffled = sorted(self.shuffle) != self.shuffle

self.is_transpose = dims_are_shuffled and not augment and not drop
self.is_squeeze = drop and not dims_are_shuffled and not augment
self.is_expand_dims = augment and not dims_are_shuffled and not drop
self.is_left_expand_dims = self.is_expand_dims and (
input_ndim == 0 or new_order[-input_ndim:] == list(range(input_ndim))
)
self.is_right_expand_dims = self.augment and new_order[:input_ndim] == list(
range(input_ndim)
)
self.is_right_expand_dims = self.is_expand_dims and new_order[
:input_ndim
] == list(range(input_ndim))

if self.inplace:
self.view_map = {0: [0]}
Expand Down Expand Up @@ -215,16 +220,15 @@ def make_node(self, inp):
return Apply(self, [input], [output])

def __str__(self):
shuffle = sorted(self.shuffle) != self.shuffle
if self.augment and not (shuffle or self.drop):
if self.is_expand_dims:
if len(self.augment) == 1:
return f"ExpandDims{{axis={self.augment[0]}}}"
return f"ExpandDims{{axes={self.augment}}}"
if self.drop and not (self.augment or shuffle):
if self.is_squeeze:
if len(self.drop) == 1:
return f"DropDims{{axis={self.drop[0]}}}"
return f"DropDims{{axes={self.drop}}}"
if shuffle and not (self.augment or self.drop):
return f"Squeeze{{axis={self.drop[0]}}}"
return f"Squeeze{{axes={self.drop}}}"
if self.is_transpose:
return f"Transpose{{axes={self.shuffle}}}"
return f"DimShuffle{{order=[{','.join(map(str, self.new_order))}]}}"

Expand Down
Loading