fix

Author: wsmoses

src/enzyme_ad/jax/gpu.cc

@@ -37,8 +37,6 @@ struct CuFuncWrapper {
   void *func;
 };
 
-void noop(void *){};
-
 template <bool withError>
 XLA_FFI_Error *initialize(XLA_FFI_CallFrame *call_frame) {
   assert(call_frame->attrs.size == 1);
@@ -83,7 +81,7 @@ XLA_FFI_Error *initialize(XLA_FFI_CallFrame *call_frame) {
   auto *execution_state = reinterpret_cast<xla::ffi::ExecutionState *>(
       internal_api->XLA_FFI_INTERNAL_ExecutionState_Get(ctx));
   (void)execution_state->Set(xla::ffi::TypeRegistry::GetTypeId<CuFuncWrapper>(),
-                             cufunc, noop);
+                             cufunc);
 
   return nullptr;
 }

third_party/jax/workspace.bzl

@@ -30,6 +30,6 @@ def repo(extra_patches = [], override_commit = ""):
         strip_prefix = "jax-" + commit,
         urls = ["https://github.com/google/jax/archive/{commit}.tar.gz".format(commit = commit)],
         patch_cmds = JAX_PATCHES + extra_patches,
-	patches = ["//:patches/jax.patch"],
+        patches = ["//:patches/jax.patch"],
         patch_args = ["-p1"],
     )