Wrapping a shard_map function

[ASKabalan]

Hello,

I am trying to wrap a shardmapped function
I made this MWE

def impl(*params, **kwargs):
    out = sharded_impl(*params)
    return out

My goal is eventually to add some operations before and after.
With parameters like this : params ((256, 256, 256), Traced<ShapedArray(uint32[2])>with<DynamicJaxprTrace(level=1/0)>) 2 type <class 'tuple'>

The first being a static with in_specs = (None , P('axis1'))
I am getting this error

ValueError: safe_zip() argument 2 is longer than argument 1

How can I use *args to bind parameters if my args contains a tracer and a static

[mattjj]

Let's discuss this on the bug thread #22356!