sketch: setup for new key array implementation based on eltypes

Co-authored-by: Matthew Johnson <[email protected]>
khcs · Aug 19, 2022 · acb5e49 · acb5e49
1 parent 39d54bd
commit acb5e49
Show file tree

Hide file tree

Showing 4 changed files with 282 additions and 25 deletions.
diff --git a/jax/_src/dispatch.py b/jax/_src/dispatch.py
@@ -108,9 +108,14 @@ def apply_primitive(prim, *args, **params):
                                         **params)
   return compiled_fun(*args)
 
-# TODO(phawkins): update code referring to xla.apply_primitive to point here.
+# TODO(phawkins,frostig,mattjj): update code referring to
+# xla.apply_primitive to point here, or use simple_impl if that's why
+# it is using apply_primitive to begin with
 xla.apply_primitive = apply_primitive
 
+def simple_impl(prim):
+  prim.def_impl(partial(apply_primitive, prim))
+
 RuntimeToken = Any
 
 class RuntimeTokenSet(threading.local):