flatten unflatten, why

Author: andrinr

examples/simple/partial/tesseract_api.py

@@ -44,10 +44,20 @@ def apply(inputs: InputSchema) -> OutputSchema:
 
 
 
+# def abstract_eval(abstract_inputs):
+#     """Calculate output shape of apply from the shape of its inputs."""
+#     return {
+#         "b": ShapeDType(shape=(None,), dtype="float32"),
+#         "c": ShapeDType(shape=(None,), dtype="float32"),
+#     }
+
+
+
+
 def abstract_eval(abstract_inputs):
     """Calculate output shape of apply from the shape of its inputs."""
     return {
-        "b": ShapeDType(shape=(None,), dtype="float32"),
-        "c": ShapeDType(shape=(None,), dtype="float32"),
+        "b": ShapeDType(shape=(abstract_inputs.a.shape[0],), dtype="float32"),
+        "c": ShapeDType(shape=(3,), dtype="float32"),
     }
 

tesseract_jax/tesseract_compat.py

@@ -159,8 +159,8 @@ def apply(
 
         out_data = self.client.apply(inputs)
 
-        out_data = tuple(jax.tree.flatten(out_data)[0])
-        return out_data
+        out_data, output_pytreedef = jax.tree.flatten(out_data)
+        return tuple(out_data), output_pytreedef
 
     def jacobian_vector_product(
         self,

test.ipynb

@@ -2,15 +2,15 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 3,
    "id": "4317d6a5",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "(array([2., 4., 6.], dtype=float32), array([1., 1., 1.], dtype=float32))\n"
+      "[Array([2., 4., 6.], dtype=float32), Array([1., 1., 1.], dtype=float32)]\n"
      ]
     }
    ],
@@ -20,14 +20,15 @@
     "from tesseract_core import Tesseract\n",
     "from tesseract_jax import apply_tesseract\n",
     "import jax.numpy as jnp\n",
+    "import jax\n",
     "\n",
     "vectoradd = Tesseract.from_tesseract_api(\"examples/simple/partial/tesseract_api.py\")\n",
     "\n",
     "\n",
     "input_dict = {\"a\": jnp.array([1.0, 2.0, 3.0], dtype=\"float32\")}\n",
     "\n",
     "\n",
-    "outputs = apply_tesseract(vectoradd, inputs=input_dict)\n",
+    "outputs = jax.jit(apply_tesseract)(vectoradd, inputs=input_dict)\n",
     "pprint(outputs)"
    ]
   },