Plumb prim params for call discharge rule (to handle named_computation_p

sharadmv · copybara-github · commit 19d8652447de · 2026-01-12T13:26:02.000-08:00
from shardy).

PiperOrigin-RevId: 853813788
diff --git a/shardy/integrations/python/jax/mpmd/ops.py b/shardy/integrations/python/jax/mpmd/ops.py
@@ -294,7 +294,9 @@ def custom_call_transpose(params, *rest, primitive=primitive):
       pe.partial_eval_jaxpr_custom_rules[primitives.call_p]
   )
   state_discharge.register_discharge_rule(primitive)(
-      state_discharge._call_discharge_rule
+      functools.partial(
+          state_discharge._call_primitive_discharge_rule, primitive
+      )
   )
   return primitive
 

Original file line number	Diff line number	Diff line change
`@@ -294,7 +294,9 @@ def custom_call_transpose(params, *rest, primitive=primitive):`
`294`	`294`	`pe.partial_eval_jaxpr_custom_rules[primitives.call_p]`
`295`	`295`	`)`
`296`	`296`	`state_discharge.register_discharge_rule(primitive)(`
`297`		`- state_discharge._call_discharge_rule`
	`297`	`+ functools.partial(`
	`298`	`+ state_discharge._call_primitive_discharge_rule, primitive`
	`299`	`+ )`
`298`	`300`	`)`
`299`	`301`	`return primitive`
`300`	`302`