Merge pull request jax-ml#8633 from shawwn:2021-11-19/autodidax-fix-j…

…axpr-subcomp-return-type PiperOrigin-RevId: 519745476
mbmccoy · Mar 27, 2023 · af4d494 · af4d494
2 parents 10d51c7 + 2d61a5f
commit af4d494
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/docs/autodidax.ipynb b/docs/autodidax.ipynb
@@ -2038,7 +2038,7 @@
    "outputs": [],
    "source": [
     "def jaxpr_subcomp(c: xe.XlaBuilder, jaxpr: Jaxpr, args: List[xe.XlaOp]\n",
-    "                  ) -> xe.XlaOp:\n",
+    "                  ) -> List[xe.XlaOp]:\n",
     "  env: Dict[Var, xe.XlaOp] = {}\n",
     "\n",
     "  def read(x: Atom) -> xe.XlaOp:\n",

diff --git a/docs/autodidax.md b/docs/autodidax.md
@@ -1598,7 +1598,7 @@ compiled program:
 
 ```{code-cell}
 def jaxpr_subcomp(c: xe.XlaBuilder, jaxpr: Jaxpr, args: List[xe.XlaOp]
-                  ) -> xe.XlaOp:
+                  ) -> List[xe.XlaOp]:
   env: Dict[Var, xe.XlaOp] = {}
 
   def read(x: Atom) -> xe.XlaOp:

diff --git a/docs/autodidax.py b/docs/autodidax.py
@@ -1592,7 +1592,7 @@ def _xla_shape(aval: ShapedArray) -> xe.Shape:
 
 # +
 def jaxpr_subcomp(c: xe.XlaBuilder, jaxpr: Jaxpr, args: List[xe.XlaOp]
-                  ) -> xe.XlaOp:
+                  ) -> List[xe.XlaOp]:
   env: Dict[Var, xe.XlaOp] = {}
 
   def read(x: Atom) -> xe.XlaOp: