pytensor icon indicating copy to clipboard operation
pytensor copied to clipboard

Useless `dimshuffle("x","x")`

Open ricardoV94 opened this issue 10 months ago • 0 comments

Description

import pytensor
import pytensor.tensor as pt

x = pt.tensor("x", shape=(None, None))
y = x.dimshuffle("x", "x")
pytensor.function([x], y).dprint()
# DeepCopyOp [id A] 1
#  └─ DimShuffle{order=[x,x]} [id B] 0
#     └─ x [id C]

Whenever a DimShuffle only has x and it has as many entries as the inputs, it's basically a specify_shape(x, (1,) * x.type.ndim). If the dimensions were known to be length1, the DimShuffle would be removed correctly

This is a low priority one since the DimShuffle is little more expensive than the specify_shape

ricardoV94 avatar Feb 24 '25 17:02 ricardoV94