Propagate static shape in MaxAndArgmax

[covertg]

Motivation for these changes

MaxAndArgmax propagates broadcastable flags but not static shape information. Now we have full shape info. Connected to #203

>>> x = pt.max_and_argmax([[1,2,3],[4,5,6]],axis=1)
>>> x[0].type.shape, x[1].type.shape 
((2), (2,))

Implementation details

I've updated the test suite to test this in a number of places. I'd appreciate close eyes on this to make sure I did the tests right.

I also parametrized a few of the updated tests to help me identify a few cases that were failing. Currently I've left them that way since I figure that less monolithic tests are preferable.

Checklist

• Explain motivation and implementation 👆
• Make sure that the pre-commit linting/style checks pass.
• Link relevant issues, preferably in nice commit messages.
• The commits correspond to relevant logical changes. Note that if they don't, we will rewrite/rebase/squash the git history before merging.
• Are the changes covered by tests and docstrings?
• Fill out the short summary sections 👇

Major / Breaking Changes

na

New features

as stated

Bugfixes

na

Documentation

• ...

Maintenance

• ...

[codecov-commenter]

Codecov Report

Merging #206 (871177b) into main (a6e7722) will not change coverage.
The diff coverage is 100.00%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #206   +/-   ##
=======================================
  Coverage   80.06%   80.06%           
=======================================
  Files         170      170           
  Lines       45185    45185           
  Branches     9608     9608           
=======================================
  Hits        36176    36176           
  Misses       6794     6794           
  Partials     2215     2215           

Impacted Files	Coverage Δ
pytensor/tensor/math.py	90.69% <100.00%> (ø)

[michaelosthege]

LGTM, thanks @covertg !

@ricardoV94 @ferrine does one of you want to take a look and merge?

[ricardoV94]

LGTM