add transform_dag implementation to inline functions for pytato JAX array context

Author: majosm

arraycontext/impl/pytato/__init__.py

@@ -817,6 +817,9 @@ class PytatoJAXArrayContext(_BasePytatoArrayContext):
     An arraycontext that uses :mod:`pytato` to represent the thawed state of
     the arrays and compiles the expressions using
     :class:`pytato.target.python.JAXPythonTarget`.
+
+
+    .. automethod:: transform_dag
     """
 
     def __init__(self,
@@ -967,6 +970,17 @@ def compile(self, f: Callable[..., Any]) -> Callable[..., Any]:
         from .compile import LazilyJAXCompilingFunctionCaller
         return LazilyJAXCompilingFunctionCaller(self, f)
 
+    def transform_dag(self, dag: pytato.DictOfNamedArrays
+                      ) -> pytato.DictOfNamedArrays:
+        import pytato as pt
+
+        # FIXME: Having to use _verify_is_dag seems clunky, but I'm not sure how to
+        # avoid it
+        from .utils import _verify_is_dag
+        dag = _verify_is_dag(pt.tag_all_calls_to_be_inlined(dag))
+        dag = _verify_is_dag(pt.inline_calls(dag))
+        return dag
+
     def tag(self, tags: ToTagSetConvertible, array):
         def _tag(ary):
             import jax.numpy as jnp