Cannot vmap mx.convolve

[williambdean]

The PyMC team is trying to implement mlx backend for pytensor in pymc-devs/pytensor#1365

We are running into issue with mx.convolve not being able to be vectorized using vmap.

[Primitive::vmap] Not implemented for Convolution.

[awni]

The PyMC team is trying to implement mlx backend

Cool!

We should add vmap for conv.

[ParamThakkar123]

I can take on this issue, would just need some pointers on how to go forward with it. Like some references or steps

[williambdean]

I have to defer that one to @ricardoV94