Implement Dimshuffle in jax using expand_dims/squeeze#987

HarshvirSandhu wants to merge 1 commit intopymc-devs:mainpymc-devs/pytensor:mainfrom

HarshvirSandhu:jax-dimshuffleHarshvirSandhu/pytensor:jax-dimshuffleCopy head branch name to clipboard