From 1d5cf7277c20e5b179053804502761a593032ad3 Mon Sep 17 00:00:00 2001
From: flaport <floris.laporte@gmail.com>
Date: Tue, 19 Sep 2023 11:54:00 -0700
Subject: [PATCH] use register_custom_call_target in stead of
 register_cpu_custom_call_target

---
 klujax.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/klujax.py b/klujax.py
index 297a73d..937f1a5 100644
--- a/klujax.py
+++ b/klujax.py
@@ -57,7 +57,7 @@ def xla_register_cpu(primitive, cpp_fun):
     name = primitive.name.encode()
 
     def decorator(fun):
-        xla_client.register_cpu_custom_call_target(
+        xla_client.register_custom_call_target(
             name,
             cpp_fun(),
         )