add concrete eval

Author: andrinr

examples/simple/partial/tesseract_api.py

@@ -31,23 +31,23 @@ def apply(inputs: InputSchema) -> OutputSchema:
     }
 
 
-def vector_jacobian_product(
-    inputs: InputSchema,
-    vjp_inputs: set[str],
-    vjp_outputs: set[str],
-    cotangent_vector: dict[str, Any],
-):
+# def vector_jacobian_product(
+#     inputs: InputSchema,
+#     vjp_inputs: set[str],
+#     vjp_outputs: set[str],
+#     cotangent_vector: dict[str, Any],
+# ):
 
-    return {
-        "a": 2.0 * cotangent_vector["b"],
-    }
+#     return {
+#         "a": 2.0 * cotangent_vector["b"],
+#     }
 
 
 
-# def abstract_eval(abstract_inputs):
-#     """Calculate output shape of apply from the shape of its inputs."""
-#     return {
-#         "b": ShapeDType(shape=(abstract_inputs.a.shape[0],), dtype="float32"),
-#         # "c": ShapeDType(shape=(3,), 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"),
+    }
 

tesseract_jax/primitive.py

@@ -21,9 +21,9 @@
 
 tesseract_dispatch_p = extend.core.Primitive("tesseract_dispatch")
 tesseract_dispatch_p.multiple_results = True
-tesseract_dispatch_p.def_impl(
-    functools.partial(xla.apply_primitive, tesseract_dispatch_p)
-)
+# tesseract_dispatch_p.def_impl(
+#     functools.partial(xla.apply_primitive, tesseract_dispatch_p)
+# )
 
 
 class _Hashable:
@@ -58,6 +58,7 @@ def tesseract_dispatch_abstract_eval(
     client: Jaxeract,
     eval_func: str,
 ) -> tuple:
+
     """Define how to dispatch evals and pipe arguments."""
     if eval_func not in (
         "apply",
@@ -78,6 +79,7 @@ def tesseract_dispatch_abstract_eval(
     return tuple(jax.core.ShapedArray(aval.shape, aval.dtype) for aval in output_avals)
 
 
+
 def tesseract_dispatch_jvp_rule(
     in_args: tuple[ArrayLike, ...],
     tan_args: tuple[ArrayLike, ...],
@@ -190,6 +192,38 @@ def tesseract_dispatch_transpose_rule(
 ad.primitive_transposes[tesseract_dispatch_p] = tesseract_dispatch_transpose_rule
 
 
+def tesseract_dispatch(
+    *array_args: ArrayLike | ShapedArray | Any,
+    static_args: tuple[_Hashable, ...],
+    input_pytreedef: PyTreeDef,
+    output_pytreedef: PyTreeDef | None,
+    output_avals: tuple[ShapeDtypeStruct, ...] | None,
+    is_static_mask: tuple[bool, ...],
+    client: Jaxeract,
+    eval_func: str,
+) -> Any:
+    """Defines how to dispatch lowering the computation."""
+
+    def _dispatch(*args: ArrayLike) -> Any:
+        static_args_ = tuple(_unpack_hashable(arg) for arg in static_args)
+        out = getattr(client, eval_func)(
+            args,
+            static_args_,
+            input_pytreedef,
+            output_pytreedef,
+            output_avals,
+            is_static_mask,
+        )
+        if not isinstance(out, tuple):
+            out = (out,)
+        return out
+    
+    result = _dispatch(*array_args)
+
+    return result
+
+tesseract_dispatch_p.def_impl(tesseract_dispatch)
+
 def tesseract_dispatch_lowering(
     ctx: Any,
     *array_args: ArrayLike | ShapedArray | Any,
@@ -406,4 +440,4 @@ def apply_tesseract(
     #     )
 
     # Unflatten the output
-    return jax.tree.unflatten(output_pytreedef, out)
+    return out#jax.tree.unflatten(output_pytreedef, out)