[Pallas] Add support for run_state to cost estimator.

PiperOrigin-RevId: 703543961

Author: justinjfu

jax/_src/pallas/cost_estimate.py

@@ -13,14 +13,17 @@
 # limitations under the License.
 """Helper tool for automatic cost estimation."""
 import dataclasses
+import functools
 import math
 from typing import Any, Sequence
 
 import jax
+from jax._src import api_util
 from jax._src import core as jax_core
 from jax._src import custom_derivatives
 from jax._src import linear_util as lu
 from jax._src import pjit
+from jax._src.state import discharge
 from jax._src.pallas import core as pallas_core
 from jax._src.interpreters import partial_eval as pe
 from jax._src.util import safe_map
@@ -87,10 +90,9 @@ def estimate_cost(fun, *args, **kwargs) -> pallas_core.CostEstimate:
     A pallas_core.CostEstimate object containing the cost estimate.
   """
   flattened_args, treedef = jax.tree.flatten(args)
-  def _partial_fun(*flat_args):
-    return fun(*jax.tree.unflatten(treedef, flat_args), **kwargs)
-  wrapped_fun = lu.wrap_init(
-      lambda *args, **kwargs: (_partial_fun(*args, **kwargs),))
+  partial_fun = functools.partial(fun, **kwargs)
+  wrapped_fun, _ = api_util.flatten_fun_nokwargs(lu.wrap_init(partial_fun),
+                                                 treedef)
   avals = [jax_core.ShapedArray(a.shape, a.dtype) for a in flattened_args]
   jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(wrapped_fun, avals)
   estimate = cost_estimate_jaxpr(jax_core.ClosedJaxpr(jaxpr, consts))
@@ -243,3 +245,12 @@ def _custom_vjp_rule(ctx, *, fun_jaxpr: jax_core.ClosedJaxpr, **_):
       bytes_accessed=inner_cost.bytes_accessed,
   )
 register_cost_rule(custom_derivatives.custom_vjp_call_jaxpr_p, _custom_vjp_rule)
+
+def _run_state_rule(*_, jaxpr: jax_core.Jaxpr, **_2):
+  inner_cost = cost_estimate_jaxpr(pe.close_jaxpr(jaxpr))
+  return CostEstimate(
+      flops=inner_cost.flops,
+      transcendentals=inner_cost.transcendentals,
+      bytes_accessed=inner_cost.bytes_accessed,
+  )
+register_cost_rule(discharge.run_state_p, _run_state_rule)

tests/pallas/pallas_cost_estimate_test.py

@@ -19,6 +19,7 @@
 from jax._src import config
 from jax._src import test_util as jtu
 from jax._src.pallas import cost_estimate
+from jax._src.state import discharge
 
 
 config.parse_flags_with_absl()
@@ -91,5 +92,23 @@ def test_integer_pow(self, power, expected_flops_per_element):
     self.assertEqual(cost.transcendentals, 0)
     self.assertEqual(cost.bytes_accessed, 80)
 
+  def test_run_state(self):
+    def add_refs(refs):
+      x_ref, y_ref, z_ref = refs
+      x = x_ref[:]
+      y = y_ref[:]
+      z = x + y
+      z_ref[:] = z
+    input_shape = jax.ShapeDtypeStruct((100,), jnp.float32)
+    cost = cost_estimate.estimate_cost(
+        discharge.run_state(add_refs),
+        (input_shape, input_shape, input_shape))
+    self.assertEqual(cost.flops, 100)
+    self.assertEqual(cost.transcendentals, 0)
+    # TODO(justinfu): This is off by a factor of 2 because run_state
+    # has all inputs/outputs as both arguments and return values.
+    self.assertEqual(cost.bytes_accessed / 2, 3 * 4 * 100)
+
+
 if __name__ == "__main__":
   absltest.main()