Don't fail when squeezing specific non-broadcastable axis

[ricardoV94]

Description

The following snippet fails:

import pytensor.tensor as pt

x = pt.vector("x")
pt.squeeze(x, axis=0)
# ValueError: Cannot drop a non-broadcastable dimension: [False], []

The user provided us information that this dimension will be of length 1, so we could add a specify_broadcastable and only then call squeeze. I think this behavior is more intuitive, and I don't see any drawbacks.

This can only be done when the user provides specific axis, not when axis=None

pt.squeeze(pt.specify_broadcastable(x, 0), axis=0).eval({x: [0]}) 
# array(0.)
pt.squeeze(pt.specify_broadcastable(x, 0), axis=0).eval({x: [0, 1, 2]})
# AssertionError: SpecifyShape: dim 0 of input has shape 3, expected 1.

[Dhruvanshu-Joshi]

So here we'll have to first make the axis(given it is not None) broadcastable manually in the pt.squeeze function itself is what is to be done?

[ricardoV94]

Yeah, but only if it's not already broadcastable (in which case it would be useless)