add transform_dag implementation to pytato JAX array context in order to inline functions

Author: majosm

arraycontext/impl/pytato/__init__.py

@@ -852,6 +852,8 @@ 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,
@@ -984,6 +986,13 @@ 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
+        dag = pt.tag_all_calls_to_be_inlined(dag)
+        dag = pt.inline_calls(dag)
+        return dag
+
     @override
     def tag(self,
                 tags: ToTagSetConvertible,