Merge pull request jax-ml#24370 from dfm:ffi-call-to-callable

PiperOrigin-RevId: 688188390

Author: Google-ML-Automation

docs/ffi.ipynb

@@ -287,27 +287,27 @@
     "  if x.dtype != jnp.float32:\n",
     "    raise ValueError(\"Only the float32 dtype is implemented by rms_norm\")\n",
     "\n",
-    "  # In this case, the output of our FFI function is just a single array with the\n",
-    "  # same shape and dtype as the input. We discuss a case with a more interesting\n",
-    "  # output type below.\n",
-    "  out_type = jax.ShapeDtypeStruct(x.shape, x.dtype)\n",
-    "\n",
-    "  return jex.ffi.ffi_call(\n",
+    "  call = jex.ffi.ffi_call(\n",
     "    # The target name must be the same string as we used to register the target\n",
     "    # above in `register_custom_call_target`\n",
     "    \"rms_norm\",\n",
-    "    out_type,\n",
-    "    x,\n",
-    "    # Note that here we're use `numpy` (not `jax.numpy`) to specify a dtype for\n",
-    "    # the attribute `eps`. Our FFI function expects this to have the C++ `float`\n",
-    "    # type (which corresponds to numpy's `float32` type), and it must be a\n",
-    "    # static parameter (i.e. not a JAX array).\n",
-    "    eps=np.float32(eps),\n",
+    "\n",
+    "    # In this case, the output of our FFI function is just a single array with\n",
+    "    # the same shape and dtype as the input. We discuss a case with a more\n",
+    "    # interesting output type below.\n",
+    "    jax.ShapeDtypeStruct(x.shape, x.dtype),\n",
+    "\n",
     "    # The `vmap_method` parameter controls this function's behavior under `vmap`\n",
     "    # as discussed below.\n",
     "    vmap_method=\"broadcast_fullrank\",\n",
     "  )\n",
     "\n",
+    "  # Note that here we're use `numpy` (not `jax.numpy`) to specify a dtype for\n",
+    "  # the attribute `eps`. Our FFI function expects this to have the C++ `float`\n",
+    "  # type (which corresponds to numpy's `float32` type), and it must be a\n",
+    "  # static parameter (i.e. not a JAX array).\n",
+    "  return call(x, eps=np.float32(eps))\n",
+    "\n",
     "\n",
     "# Test that this gives the same result as our reference implementation\n",
     "x = jnp.linspace(-0.5, 0.5, 15).reshape((3, 5))\n",
@@ -403,10 +403,8 @@
     "  return jex.ffi.ffi_call(\n",
     "    \"rms_norm\",\n",
     "    jax.ShapeDtypeStruct(x.shape, x.dtype),\n",
-    "    x,\n",
-    "    eps=np.float32(eps),\n",
     "    vmap_method=\"sequential\",\n",
-    "  )\n",
+    "  )(x, eps=np.float32(eps))\n",
     "\n",
     "\n",
     "jax.make_jaxpr(jax.vmap(rms_norm_sequential))(x)"
@@ -462,10 +460,8 @@
     "      jax.ShapeDtypeStruct(x.shape, x.dtype),\n",
     "      jax.ShapeDtypeStruct(x.shape[:-1], x.dtype),\n",
     "    ),\n",
-    "    x,\n",
-    "    eps=np.float32(eps),\n",
     "    vmap_method=\"broadcast_fullrank\",\n",
-    "  )\n",
+    "  )(x, eps=np.float32(eps))\n",
     "  return y, (res, x)\n",
     "\n",
     "\n",
@@ -478,11 +474,8 @@
     "    jex.ffi.ffi_call(\n",
     "      \"rms_norm_bwd\",\n",
     "      jax.ShapeDtypeStruct(ct.shape, ct.dtype),\n",
-    "      res,\n",
-    "      x,\n",
-    "      ct,\n",
-    "    vmap_method=\"broadcast_fullrank\",\n",
-    "    ),\n",
+    "      vmap_method=\"broadcast_fullrank\",\n",
+    "    )(res, x, ct),\n",
     "  )\n",
     "\n",
     "\n",
@@ -569,10 +562,8 @@
     "    return lambda x: jex.ffi.ffi_call(\n",
     "      target_name,\n",
     "      out_type,\n",
-    "      x,\n",
-    "      eps=np.float32(eps),\n",
     "      vmap_method=\"broadcast_fullrank\",\n",
-    "    )\n",
+    "    )(x, eps=np.float32(eps))\n",
     "\n",
     "  return jax.lax.platform_dependent(x, cpu=impl(\"rms_norm\"), cuda=impl(\"rms_norm_cuda\"))\n",
     "\n",

docs/ffi.md

@@ -248,27 +248,27 @@ def rms_norm(x, eps=1e-5):
   if x.dtype != jnp.float32:
     raise ValueError("Only the float32 dtype is implemented by rms_norm")
 
-  # In this case, the output of our FFI function is just a single array with the
-  # same shape and dtype as the input. We discuss a case with a more interesting
-  # output type below.
-  out_type = jax.ShapeDtypeStruct(x.shape, x.dtype)
-
-  return jex.ffi.ffi_call(
+  call = jex.ffi.ffi_call(
     # The target name must be the same string as we used to register the target
     # above in `register_custom_call_target`
     "rms_norm",
-    out_type,
-    x,
-    # Note that here we're use `numpy` (not `jax.numpy`) to specify a dtype for
-    # the attribute `eps`. Our FFI function expects this to have the C++ `float`
-    # type (which corresponds to numpy's `float32` type), and it must be a
-    # static parameter (i.e. not a JAX array).
-    eps=np.float32(eps),
+
+    # In this case, the output of our FFI function is just a single array with
+    # the same shape and dtype as the input. We discuss a case with a more
+    # interesting output type below.
+    jax.ShapeDtypeStruct(x.shape, x.dtype),
+
     # The `vmap_method` parameter controls this function's behavior under `vmap`
     # as discussed below.
     vmap_method="broadcast_fullrank",
   )
 
+  # Note that here we're use `numpy` (not `jax.numpy`) to specify a dtype for
+  # the attribute `eps`. Our FFI function expects this to have the C++ `float`
+  # type (which corresponds to numpy's `float32` type), and it must be a
+  # static parameter (i.e. not a JAX array).
+  return call(x, eps=np.float32(eps))
+
 
 # Test that this gives the same result as our reference implementation
 x = jnp.linspace(-0.5, 0.5, 15).reshape((3, 5))
@@ -334,10 +334,8 @@ def rms_norm_sequential(x, eps=1e-5):
   return jex.ffi.ffi_call(
     "rms_norm",
     jax.ShapeDtypeStruct(x.shape, x.dtype),
-    x,
-    eps=np.float32(eps),
     vmap_method="sequential",
-  )
+  )(x, eps=np.float32(eps))
 
 
 jax.make_jaxpr(jax.vmap(rms_norm_sequential))(x)
@@ -380,10 +378,8 @@ def rms_norm_fwd(x, eps=1e-5):
       jax.ShapeDtypeStruct(x.shape, x.dtype),
       jax.ShapeDtypeStruct(x.shape[:-1], x.dtype),
     ),
-    x,
-    eps=np.float32(eps),
     vmap_method="broadcast_fullrank",
-  )
+  )(x, eps=np.float32(eps))
   return y, (res, x)
 
 
@@ -396,11 +392,8 @@ def rms_norm_bwd(eps, res, ct):
     jex.ffi.ffi_call(
       "rms_norm_bwd",
       jax.ShapeDtypeStruct(ct.shape, ct.dtype),
-      res,
-      x,
-      ct,
-    vmap_method="broadcast_fullrank",
-    ),
+      vmap_method="broadcast_fullrank",
+    )(res, x, ct),
   )
 
 
@@ -477,10 +470,8 @@ def rms_norm_cross_platform(x, eps=1e-5):
     return lambda x: jex.ffi.ffi_call(
       target_name,
       out_type,
-      x,
-      eps=np.float32(eps),
       vmap_method="broadcast_fullrank",
-    )
+    )(x, eps=np.float32(eps))
 
   return jax.lax.platform_dependent(x, cpu=impl("rms_norm"), cuda=impl("rms_norm_cuda"))
 

examples/ffi/src/jax_ffi_example/attrs.py

@@ -35,13 +35,11 @@ def array_attr(num: int):
   return jex.ffi.ffi_call(
       "array_attr",
       jax.ShapeDtypeStruct((), np.int32),
-      array=np.arange(num, dtype=np.int32),
-  )
+  )(array=np.arange(num, dtype=np.int32))
 
 
 def dictionary_attr(**kwargs):
   return jex.ffi.ffi_call(
       "dictionary_attr",
       (jax.ShapeDtypeStruct((), np.int32), jax.ShapeDtypeStruct((), np.int32)),
-      **kwargs,
-  )
+  )(**kwargs)

examples/ffi/src/jax_ffi_example/rms_norm.py

@@ -49,19 +49,17 @@ def rms_norm(x, eps=1e-5):
   # same shape and dtype as the input.
   out_type = jax.ShapeDtypeStruct(x.shape, x.dtype)
 
+  # Note that here we're use `numpy` (not `jax.numpy`) to specify a dtype for
+  # the attribute `eps`. Our FFI function expects this to have the C++ `float`
+  # type (which corresponds to numpy's `float32` type), and it must be a
+  # static parameter (i.e. not a JAX array).
   return jex.ffi.ffi_call(
     # The target name must be the same string as we used to register the target
     # above in `register_ffi_target`
     "rms_norm",
     out_type,
-    x,
-    # Note that here we're use `numpy` (not `jax.numpy`) to specify a dtype for
-    # the attribute `eps`. Our FFI function expects this to have the C++ `float`
-    # type (which corresponds to numpy's `float32` type), and it must be a
-    # static parameter (i.e. not a JAX array).
-    eps=np.float32(eps),
     vmap_method="broadcast_fullrank",
-  )
+  )(x, eps=np.float32(eps))
 
 
 def rms_norm_fwd(x, eps=1e-5):
@@ -71,10 +69,8 @@ def rms_norm_fwd(x, eps=1e-5):
       jax.ShapeDtypeStruct(x.shape, x.dtype),
       jax.ShapeDtypeStruct(x.shape[:-1], x.dtype),
     ),
-    x,
-    eps=np.float32(eps),
     vmap_method="broadcast_fullrank",
-  )
+  )(x, eps=np.float32(eps))
   return y, (res, x)
 
 
@@ -87,11 +83,8 @@ def rms_norm_bwd(eps, res, ct):
     jex.ffi.ffi_call(
       "rms_norm_bwd",
       jax.ShapeDtypeStruct(ct.shape, ct.dtype),
-      res,
-      x,
-      ct,
       vmap_method="broadcast_fullrank",
-    ),
+    )(res, x, ct),
   )