Add concrete_mesh to reshard_p in pallas lowering rules which was missing

PiperOrigin-RevId: 845766226

Author: yashk2810

jax/_src/pallas/mosaic/lowering.py

@@ -3393,7 +3393,8 @@ def _pjit_lowering_rule(ctx: LoweringRuleContext, *args, jaxpr, **_):
 
 
 @register_lowering_rule(pjit.reshard_p)
-def _reshard_lowering_rule(ctx: LoweringRuleContext, x, dst_sharding):
+def _reshard_lowering_rule(ctx: LoweringRuleContext, x, *, dst_sharding,
+                           concrete_mesh):
   return x
 
 

jax/_src/pallas/triton/lowering.py

@@ -2537,7 +2537,7 @@ def _pjit_lowering_rule(ctx: LoweringRuleContext, *args, jaxpr, **_):
 
 
 @register_lowering(pjit.reshard_p)
-def _reshard_lowering_rule(ctx, x, dst_sharding):
+def _reshard_lowering_rule(ctx, x, *, dst_sharding, concrete_mesh):
   return x