Foreign function interface (FFI)#
This tutorial requires JAX v0.4.31 or newer. The later sections that add custom transformation rules also use the experimental jax.experimental.hijax API, which requires a more recent version of JAX.
While a wide range of numerical operations can be easily and efficiently implemented using JAX’s built in jax.numpy and jax.lax interfaces, it can sometimes be useful to explicitly call out to external compiled libraries via a “foreign function interface” (FFI).
This can be particularly useful when particular operations have been previously implemented in an optimized C or CUDA library, and it would be non-trivial to reimplement these computations directly using JAX, but it can also be useful for optimizing runtime or memory performance of JAX programs.
That being said, the FFI should typically be considered a last resort option because the XLA compiler that sits in the backend, or the Pallas kernel language, which provides lower level control, typically produce performant code with a lower development and maintenance cost.
One point that should be taken into account when considering use of the FFI is that JAX doesn’t automatically know how to differentiate through foreign functions. This means that if you want to use JAX’s autodifferentiation capabilities alongside a foreign function, you’ll also need to provide an implementation of the relevant differentiation rules. We will discuss some possible approaches below, but it is important to call this limitation out right from the start!
JAX’s FFI support is provided in two parts:
A header-only C++ library from XLA which is packaged as part of JAX as of v0.4.29 or available from the openxla/xla project, and
A Python front end, available in the
jax.ffisubmodule.
In this tutorial we demonstrate the use of both of these components using a simple example, and then go on to discuss some lower-level extensions for more complicated use cases. We start by presenting the FFI on CPU, and discuss generalizations to GPU or multi-device environments below.
The end-to-end code for this example and some other more advanced use cases can be found in the JAX FFI examples project on GitHub at examples/ffi in the JAX repository.
Because we demonstrate how FFI calls can be sharded at the end of this tutorial, let’s first set up our environment to be treated by JAX as having multiple CPUs:
# ruff: noqa
import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=4"
A simple example#
To demonstrate the use of the FFI interface, we will implement a simple “root-mean-square (RMS)” normalization function. RMS normalization takes an array \(x\) with shape \((N,)\) and returns
where \(\epsilon\) is a tuning parameter used for numerical stability.
This is a somewhat silly example, because it can be easily implemented using JAX as follows:
import jax
import jax.numpy as jnp
def rms_norm_ref(x, eps=1e-5):
scale = jnp.sqrt(jnp.mean(jnp.square(x), axis=-1, keepdims=True) + eps)
return x / scale
But, it’s just non-trivial enough to be useful for demonstrating some key details of the FFI, while still being straightforward to understand. We will use this reference implementation to test our FFI version below.
Backend code#
To begin with, we need an implementation of RMS normalization in C++ that we will expose using the FFI. This isn’t meant to be particularly performant, but you could imagine that if you had some new better implementation of RMS normalization in a C++ library, it might have an interface like the following. So, here’s a simple implementation of RMS normalization in C++:
#include <cmath>
#include <cstdint>
float ComputeRmsNorm(float eps, int64_t size, const float *x, float *y) {
float sm = 0.0f;
for (int64_t n = 0; n < size; ++n) {
sm += x[n] * x[n];
}
float scale = 1.0f / std::sqrt(sm / float(size) + eps);
for (int64_t n = 0; n < size; ++n) {
y[n] = x[n] * scale;
}
return scale;
}
and, for our example, this is the function that we want to expose to JAX via the FFI.
C++ interface#
To expose our library function to JAX and XLA, we need to write a thin wrapper using the APIs provided by the header-only library in the xla/ffi/api directory of the XLA project.
For more information about this interface, take a look at the XLA custom call documentation.
The full source listing can be downloaded here, but the key implementation details are reproduced here (simplified to support only the float32 dtype; the full source instead uses ffi::AnyBuffer to support several dtypes, as discussed under Advanced topics below):
#include <functional>
#include <numeric>
#include <utility>
#include "xla/ffi/api/c_api.h"
#include "xla/ffi/api/ffi.h"
namespace ffi = xla::ffi;
// A helper function for extracting the relevant dimensions from `ffi::Buffer`s.
// In this example, we treat all leading dimensions as batch dimensions, so this
// function returns the total number of elements in the buffer, and the size of
// the last dimension.
template <ffi::DataType T>
std::pair<int64_t, int64_t> GetDims(const ffi::Buffer<T> &buffer) {
auto dims = buffer.dimensions();
if (dims.size() == 0) {
return std::make_pair(0, 0);
}
return std::make_pair(buffer.element_count(), dims.back());
}
// A wrapper function providing the interface between the XLA FFI call and our
// library function `ComputeRmsNorm` above. This function handles the batch
// dimensions by calling `ComputeRmsNorm` within a loop.
ffi::Error RmsNormImpl(float eps, ffi::Buffer<ffi::F32> x,
ffi::ResultBuffer<ffi::F32> y) {
auto [totalSize, lastDim] = GetDims(x);
if (lastDim == 0) {
return ffi::Error::InvalidArgument("RmsNorm input must be an array");
}
for (int64_t n = 0; n < totalSize; n += lastDim) {
ComputeRmsNorm(eps, lastDim, &(x.typed_data()[n]), &(y->typed_data()[n]));
}
return ffi::Error::Success();
}
// Wrap `RmsNormImpl` and specify the interface to XLA. If you need to declare
// this handler in a header, you can use the `XLA_FFI_DECLARE_HANDLER_SYMBOL`
// macro: `XLA_FFI_DECLARE_HANDLER_SYMBOL(RmsNorm)`.
XLA_FFI_DEFINE_HANDLER_SYMBOL(
RmsNorm, RmsNormImpl,
ffi::Ffi::Bind()
.Attr<float>("eps")
.Arg<ffi::Buffer<ffi::F32>>() // x
.Ret<ffi::Buffer<ffi::F32>>() // y
);
Starting at the bottom, we’re using the XLA-provided macro XLA_FFI_DEFINE_HANDLER_SYMBOL to generate some boilerplate which will expand into a function called RmsNorm with the appropriate signature.
But, the important stuff here is all in the call to ffi::Ffi::Bind(), where we define the input and output types, and the types of any parameters.
Then, in RmsNormImpl, we accept ffi::Buffer arguments which include information about the buffer shape, and pointers to the underlying data.
In this implementation, we treat all leading dimensions of the buffer as batch dimensions, and perform RMS normalization over the last axis.
GetDims is a helper function providing support for this batching behavior.
We discuss this batching behavior in more detail below, but the general idea is that it can be useful to transparently handle batching in the left-most dimensions of the input arguments.
In this case, we treat all but the last axis as batch dimensions, but other foreign functions may require a different number of non-batch dimensions.
Building and registering an FFI handler#
Now that we have our minimal FFI wrapper implemented, we need to expose this function (RmsNorm) to Python.
In this tutorial, we compile RmsNorm into a shared library and load it using ctypes, but another common pattern is to use nanobind or pybind11 as discussed below.
To compile the shared library, we’re using CMake here, but you should be able to use your favorite build system without too much trouble.
!cmake -DCMAKE_BUILD_TYPE=Release -B ffi/_build ffi
!cmake --build ffi/_build
!cmake --install ffi/_build
With this compiled library in hand, we now need to register this handler with XLA via the register_ffi_target() function.
This function expects our handler (a function pointer to the C++ function RmsNorm) to be wrapped in a PyCapsule.
JAX provides a helper function pycapsule() to help with this:
import ctypes
from pathlib import Path
path = next(Path("ffi").glob("librms_norm*"))
rms_norm_lib = ctypes.cdll.LoadLibrary(path)
jax.ffi.register_ffi_target(
"rms_norm", jax.ffi.pycapsule(rms_norm_lib.RmsNorm), platform="cpu")
Tip
If you’re familiar with the legacy “custom call” API, it’s worth noting that you can also use register_ffi_target() to register a custom call target by manually specifying the keyword argument api_version=0. The default api_version for register_ffi_target() is 1, the new “typed” FFI API that we’re using here.
An alternative approach: A common alternative pattern for exposing handlers to Python is to use nanobind or pybind11 to define a tiny Python extension which can be imported. For our example here, the nanobind code would be:
#include <type_traits>
#include "nanobind/nanobind.h"
#include "xla/ffi/api/c_api.h"
namespace nb = nanobind;
template <typename T>
nb::capsule EncapsulateFfiCall(T *fn) {
// This check is optional, but it can be helpful for avoiding invalid handlers.
static_assert(std::is_invocable_r_v<XLA_FFI_Error *, T, XLA_FFI_CallFrame *>,
"Encapsulated function must be and XLA FFI handler");
return nb::capsule(reinterpret_cast<void *>(fn));
}
NB_MODULE(rms_norm, m) {
m.def("rms_norm", []() { return EncapsulateFfiCall(RmsNorm); });
}
Then, in Python we can register this handler using:
# Assuming that we compiled a nanobind extension called `rms_norm`:
import rms_norm as rms_norm_lib
jax.ffi.register_ffi_target("rms_norm", rms_norm_lib.rms_norm(), platform="cpu")
Frontend code#
Now that we have registered our FFI handler, it is straightforward to call our C++ library from JAX using the ffi_call() function:
import numpy as np
def rms_norm(x, eps=1e-5):
# We only implemented the `float32` version of this function, so we start by
# checking the dtype. This check isn't strictly necessary because type
# checking is also performed by the FFI when decoding input and output
# buffers, but it can be useful to check types in Python to raise more
# informative errors.
if x.dtype != jnp.float32:
raise ValueError("Only the float32 dtype is implemented by rms_norm")
call = jax.ffi.ffi_call(
# The target name must be the same string as we used to register the target
# above with `jax.ffi.register_ffi_target`
"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.
jax.ShapeDtypeStruct.like(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).
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, 32).reshape((8, 4))
np.testing.assert_allclose(rms_norm(x), rms_norm_ref(x), rtol=1e-5)
This code cell includes a lot of inline comments which should explain most of what is happening here, but there are a few points that are worth explicitly highlighting.
Most of the heavy lifting here is done by the ffi_call() function, which tells JAX how to call the foreign function for a particular set of inputs.
It’s important to note that the first argument to ffi_call() must be a string that matches the target name that we used when calling register_ffi_target() above.
Any attributes (defined using Attr in the C++ wrapper above) should be passed as keyword arguments to ffi_call().
Note that we explicitly cast eps to np.float32 because our FFI library expects a C float, and we can’t use jax.numpy here, because these parameters must be static arguments.
Supporting transformations like vmap and grad#
The rms_norm function above works whenever we evaluate it directly, but as written it doesn’t yet support JAX’s function transformations.
As far as JAX is concerned, an ffi_call() is an opaque black box: JAX can’t look inside it to work out how it should behave under vmap(), or how to differentiate it.
So, for example, trying to differentiate rms_norm as defined above would fail.
To teach JAX how to transform our foreign function, we wrap it in a HiJAX primitive: a custom JAX operation defined by subclassing VJPHiPrimitive from the experimental jax.experimental.hijax module.
On the primitive, we define a handful of methods:
expandimplements the operation in terms of other (“lojax”) JAX operations. Here, that’s just a call toffi_call(), and it is what runs when the primitive isn’t being transformed.vjp_fwdandvjp_bwd_retvaltogether define the reverse-mode automatic differentiation rule (used bygrad(),vjp(), and friends).jvpdefines the forward-mode rule (used byjvp()). Unlikecustom_vjp(), a single HiJAX primitive can supply both forward- and reverse-mode AD rules.batchdefines thevmap()rule.
Note
HiJAX is a new and experimental API. The details of the interface may change in future releases of JAX.
To support differentiation, we use two additional FFI targets for the forward and backward passes:
rms_norm_fwdreturns two outputs: (a) the “primal” result, and (b) the “residuals” which are saved for use on the backwards pass, andrms_norm_bwdtakes the residuals and the output co-tangents, and returns the input co-tangents.
We won’t get into the details of the RMS normalization backwards pass, but take a look at the C++ source code to see how these functions are implemented on the back end.
Note that the residual computed by rms_norm_fwd has a different shape than the primal output (it stores one value per row), so its ffi_call() returns two outputs with different shapes.
We’ll access these two FFI targets through small helper functions, rms_norm_fwd and rms_norm_bwd, and call those from the primitive’s AD rules:
from jax.experimental.hijax import VJPHiPrimitive
jax.ffi.register_ffi_target(
"rms_norm_fwd", jax.ffi.pycapsule(rms_norm_lib.RmsNormFwd), platform="cpu")
jax.ffi.register_ffi_target(
"rms_norm_bwd", jax.ffi.pycapsule(rms_norm_lib.RmsNormBwd), platform="cpu")
def rms_norm_fwd(x, eps):
# Returns the primal output `y` and the residual `res` (one value per row, so
# it has a different shape than `y`) that the backward pass needs.
out_type = (jax.typeof(x), jax.ShapeDtypeStruct(x.shape[:-1], x.dtype))
return jax.ffi.ffi_call("rms_norm_fwd", out_type)(x, eps=eps)
def rms_norm_bwd(res, x, ct):
# Maps the output co-tangent `ct` back to the input co-tangent.
return jax.ffi.ffi_call("rms_norm_bwd", jax.typeof(x))(res, x, ct)
class RMSNorm(VJPHiPrimitive):
def __init__(self, aval, eps):
if aval.dtype != jnp.float32:
raise ValueError("Only the float32 dtype is implemented by rms_norm")
self.in_avals = (aval,)
self.out_aval = aval
self.params = dict(eps=np.float32(eps))
super().__init__()
def expand(self, x):
# The plain implementation: a single FFI call, used when we aren't
# transforming.
return jax.ffi.ffi_call("rms_norm", self.out_aval)(x, eps=self.eps)
def vjp_fwd(self, nzs_in, x):
# `nzs_in` flags which inputs have nonzero tangents; we don't need it here.
y, res = rms_norm_fwd(x, self.eps)
return y, (res, x)
def vjp_bwd_retval(self, res, ct):
res, x = res
return (rms_norm_bwd(res, x, ct),)
def jvp(self, primals, tangents):
(x,), (dx,) = primals, tangents
y, res = rms_norm_fwd(x, self.eps)
# RMS normalization has a symmetric Jacobian, so the same `rms_norm_bwd`
# kernel serves both the reverse-mode and the forward-mode rules.
return y, rms_norm_bwd(res, x, dx)
def batch(self, axis_data, args, dims):
# Our handler treats all leading axes as batch dimensions, so to support
# `vmap` we move the mapped axis to the front and reapply `rms_norm`.
x, = args
bdim, = dims
return rms_norm(jnp.moveaxis(x, bdim, 0), eps=self.eps), 0
def rms_norm(x, eps=1e-5):
return RMSNorm(jax.typeof(x), eps)(x)
Now rms_norm evaluates to the same values as before, but it also differentiates (in both forward and reverse mode) and batches:
x = jnp.linspace(-0.5, 0.5, 32).reshape((8, 4))
np.testing.assert_allclose(rms_norm(x), rms_norm_ref(x), rtol=1e-5)
np.testing.assert_allclose(jax.jit(rms_norm)(x), rms_norm_ref(x), rtol=1e-5)
# Reverse-mode AD matches the reference implementation
ct_y = jnp.ones_like(x)
np.testing.assert_allclose(
jax.vjp(rms_norm, x)[1](ct_y), jax.vjp(rms_norm_ref, x)[1](ct_y), rtol=1e-5)
# Forward-mode AD does too
dx = jnp.ones_like(x)
np.testing.assert_allclose(
jax.jvp(rms_norm, (x,), (dx,))[1],
jax.jvp(rms_norm_ref, (x,), (dx,))[1], rtol=1e-4, atol=1e-5)
# As does `vmap`
xs = jnp.linspace(-0.5, 0.5, 96).reshape((3, 8, 4))
np.testing.assert_allclose(jax.vmap(rms_norm)(xs), jax.vmap(rms_norm_ref)(xs), rtol=1e-5)
Composing transformations#
Each of grad(), jvp(), and vmap() works on its own, but composing them doesn’t work yet.
For example, vmap() of grad() (a batched gradient) currently fails with an error like Batching rule for 'ffi_call' not implemented.
The reason is that our AD rules call the rms_norm_fwd and rms_norm_bwd helpers, which call ffi_call() directly, and a raw ffi_call() has no batching rule of its own.
So when vmap() tries to push through the differentiation rule, it eventually reaches an un-batchable ffi_call() and gives up.
The fix is to give those FFI calls a batch rule, by wrapping each one in its own small HiJAX primitive (exactly as we did for rms_norm).
Because RMSNorm’s rules already call the rms_norm_fwd and rms_norm_bwd helpers by name, we only need to replace those two helpers; RMSNorm itself doesn’t change:
class RMSNormFwd(VJPHiPrimitive):
def __init__(self, aval, eps):
self.in_avals = (aval,)
self.out_aval = (aval, aval.update(shape=aval.shape[:-1])) # y, res
self.params = dict(eps=np.float32(eps))
super().__init__()
def expand(self, x):
return jax.ffi.ffi_call("rms_norm_fwd", self.out_aval)(x, eps=self.eps)
def batch(self, axis_data, args, dims):
x, = args
bdim, = dims
return rms_norm_fwd(jnp.moveaxis(x, bdim, 0), self.eps), (0, 0)
def rms_norm_fwd(x, eps):
return RMSNormFwd(jax.typeof(x), eps)(x)
class RMSNormBwd(VJPHiPrimitive):
def __init__(self, res_aval, x_aval, ct_aval):
self.in_avals = (res_aval, x_aval, ct_aval)
self.out_aval = x_aval
self.params = {}
super().__init__()
def expand(self, res, x, ct):
return jax.ffi.ffi_call("rms_norm_bwd", self.out_aval)(res, x, ct)
def batch(self, axis_data, args, dims):
# Move each mapped axis to the front; an unmapped (`None`) input is constant
# across the batch, so broadcast it along a new leading axis.
def to_front(a, d):
if d is None:
return jnp.broadcast_to(a, (axis_data.size, *a.shape))
return jnp.moveaxis(a, d, 0)
res, x, ct = (to_front(a, d) for a, d in zip(args, dims))
return rms_norm_bwd(res, x, ct), 0
def rms_norm_bwd(res, x, ct):
return RMSNormBwd(jax.typeof(res), jax.typeof(x), jax.typeof(ct))(res, x, ct)
With the helpers now batchable, the transformations compose. Here are vmap of grad and vmap of jvp, neither of which worked above:
np.testing.assert_allclose(
jax.vmap(jax.grad(lambda x: rms_norm(x).sum()))(xs),
jax.vmap(jax.grad(lambda x: rms_norm_ref(x).sum()))(xs), rtol=1e-4, atol=1e-5)
np.testing.assert_allclose(
jax.vmap(lambda x, t: jax.jvp(rms_norm, (x,), (t,))[1])(xs, jnp.ones_like(xs)),
jax.vmap(lambda x, t: jax.jvp(rms_norm_ref, (x,), (t,))[1])(xs, jnp.ones_like(xs)),
rtol=1e-4, atol=1e-5)
Note that this primitive supports only first-order AD. Higher-order AD (such as grad() of grad()) isn’t available here, because the rms_norm_bwd primitive has no differentiation rule of its own (its expand is an opaque FFI call).
FFI calls on a GPU#
So far, we have been interfacing only with foreign functions running on the CPU, but JAX’s FFI also supports calls to GPU code. Since this documentation page is automatically generated on a machine without access to a GPU, we can’t execute any GPU-specific examples here, but we will go over the key points.
When defining our FFI wrapper for CPU, the function signature that we used was:
ffi::Error RmsNormImpl(float eps, ffi::Buffer<ffi::F32> x,
ffi::ResultBuffer<ffi::F32> y)
To update this to interface with a CUDA kernel, this signature becomes:
ffi::Error RmsNormImpl(cudaStream_t stream, float eps,
ffi::Buffer<ffi::F32> x,
ffi::ResultBuffer<ffi::F32> y)
And the handler definition is updated to include a Ctx in its binding:
XLA_FFI_DEFINE_HANDLER(
RmsNorm, RmsNormImpl,
ffi::Ffi::Bind()
.Ctx<ffi::PlatformStream<cudaStream_t>>()
.Attr<float>("eps")
.Arg<ffi::Buffer<ffi::F32>>() // x
.Ret<ffi::Buffer<ffi::F32>>() // y
);
Then, the RmsNormImpl can use the CUDA stream to launch CUDA kernels.
On the front end, the registration code would be updated to specify the appropriate platform:
jax.ffi.register_ffi_target(
"rms_norm_cuda", rms_norm_lib_cuda.rms_norm(), platform="CUDA"
)
Supporting multiple platforms#
Suppose we have registered both the CPU target rms_norm and the GPU target rms_norm_cuda.
To support running our function on either platform, we can choose the FFI target inside expand, based on the platform that the computation will run on.
In explicit sharding mode the active mesh carries an abstract description of the target device, whose platform we can query to pick the right target name. Note that this returns the lowercase XLA platform name (such as "cuda"), while the platform argument to register_ffi_target() is case-insensitive, so the "CUDA" registration above matches the "cuda" lookup here:
class RMSNorm(VJPHiPrimitive):
def __init__(self, aval, eps):
if aval.dtype != jnp.float32:
raise ValueError("Only the float32 dtype is implemented by rms_norm")
self.in_avals = (aval,)
self.out_aval = aval
self.params = dict(eps=np.float32(eps))
super().__init__()
def expand(self, x):
platform = jax.sharding.get_abstract_mesh().abstract_device.platform
name = {"cpu": "rms_norm", "cuda": "rms_norm_cuda"}[platform]
body = lambda x: jax.ffi.ffi_call(name, jax.typeof(x))(x, eps=self.eps)
return jax.shard_map(body, out_specs=self.out_aval.sharding.spec)(x)
def rms_norm(x, eps=1e-5):
return RMSNorm(jax.typeof(x), eps)(x)
# On this CPU-only machine, the `cpu` target is selected:
np.testing.assert_allclose(rms_norm(x), rms_norm_ref(x), rtol=1e-5)
Even without a GPU, we can confirm that the cuda target would be selected when lowering for a GPU.
To do this, we lower (but don’t compile or run) the function under an abstract mesh that stands in for a GPU device:
gpu_device = jax.sharding.AbstractDevice("gpu", 1, "cuda")
gpu_mesh = jax.sharding.AbstractMesh(
(4,), ("x",), (jax.sharding.AxisType.Explicit,), abstract_device=gpu_device)
with jax.sharding.use_abstract_mesh(gpu_mesh):
hlo = jax.jit(rms_norm).lower(x).as_text()
print(hlo)
module @jit_rms_norm attributes {mhlo.num_partitions = 4 : i32, mhlo.num_replicas = 1 : i32} {
sdy.mesh @mesh = <["x"=4]> {stablehlo.mesh = {axes = [{name = "x", size = 4 : i64}]}}
func.func public @main(%arg0: tensor<8x4xf32>) -> (tensor<8x4xf32> {jax.result_info = "result", sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>}) {
%0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{}, {}]>] out_shardings=[<@mesh, [{}, {}]>] manual_axes={"x"} (%arg1: tensor<8x4xf32>) {
%1 = stablehlo.custom_call @rms_norm_cuda(%arg1) {backend_config = "", mhlo.backend_config = {eps = 9.99999974E-6 : f32}, operand_layouts = [dense<[1, 0]> : tensor<2xindex>], result_layouts = [dense<[1, 0]> : tensor<2xindex>]} : (tensor<8x4xf32>) -> tensor<8x4xf32>
sdy.return %1 : tensor<8x4xf32>
} : (tensor<8x4xf32>) -> tensor<8x4xf32>
return %0 : tensor<8x4xf32>
}
}
As you can see in the lowered program above, the FFI call now targets rms_norm_cuda, even though it was traced on a CPU-only machine.
The deprecated vmap_method argument#
The approach above, where vmap() support comes from a HiJAX primitive’s batch rule, is the recommended way to make a foreign function batchable.
But you may encounter older code that instead passes a vmap_method string argument directly to ffi_call().
This argument is deprecated in favor of the HiJAX primitive approach, but it’s worth knowing what it does.
When vmap_method was specified, a bare ffi_call() could be mapped under vmap() without defining a custom primitive.
For example, the batching behavior of our RMS normalization example used to be written like this:
def rms_norm(x, eps=1e-5):
return jax.ffi.ffi_call(
"rms_norm",
jax.ShapeDtypeStruct.like(x),
vmap_method="broadcast_all",
)(x, eps=np.float32(eps))
The allowed values, which mirror those of pure_callback(), were:
"sequential": apply the foreign function to one batch element at a time, in a loop (usingmap()). This is the most general but slowest option."sequential_unrolled": like"sequential", but the loop is unrolled."expand_dims": insert a new leading batch dimension and call the foreign function once; the handler is responsible for interpreting the extra dimension."broadcast_all": broadcast all of the inputs to a common batch shape and call the foreign function once (appropriate here, since our handler already treats leading axes as batch dimensions)."legacy_vectorized": a legacy mode retained for backwards compatibility.
See the pure_callback() documentation for more on the semantics of each method.
If you have code that relies on vmap_method, the recommended migration is to wrap the ffi_call() in a HiJAX primitive and implement a batch rule, as described in the transformations section above.
Advanced topics#
This tutorial covers most of the basic steps that are required to get up and running with JAX’s FFI, but advanced use cases may require more features. We will leave these topics to future tutorials, but here are some possibly useful references:
Supporting multiple dtypes: In this tutorial’s example, we restricted to only support
float32inputs and outputs, but many use cases require supporting multiple different input types. One option to handle this is to register different FFI targets for all supported input types and then use Python to select the appropriate target forjax.ffi.ffi_call()depending on the input types. But, this approach could get quickly unwieldy depending on the combinatorics of the supported cases. So it is also possible to define the C++ handler to acceptffi::AnyBufferinstead offfi::Buffer<Dtype>. Then, the input buffer will include aelement_type()method which can be used to define the appropriate dtype dispatching logic in the backend.Stateful foreign functions: It is also possible to use the FFI to wrap functions with associated state. There is a low-level example included in the XLA test suite, and a future tutorial will include more details.