Rebuilding TensorFlow Lite with Custom Ops - Integrating MediaPipe's Attention Face Mesh

Intro

While adding the attention face mesh model to the mediapipe_face_mesh package, I rebuilt the TFLite runtime binary with custom ops baked in. This is a writeup of that process.

face_landmark_with_attention.tflite refines the eyes, lips, and irises in a single inference pass and hands back 478 landmarks directly. The old path was two passes — run the 468-point mesh model, then run a separate iris model. This one does it in a single pass.

Two terms I lean on throughout.

• stock: the untouched, as-shipped state — the official build (binary) or source as distributed. I use it as the baseline against anything I rebuilt or patched with custom ops.
• kernel: the function that actually computes a single op. Not the OS kernel — the compute-kernel sense from GPU/numerics.

What a custom op is

A .tflite file is a compute graph. Input tensors flow through ops like CONV_2D → RELU → FULLY_CONNECTED → … to produce output tensors. The interpreter walks the graph and, for each op, looks up and runs the matching kernel — the function that actually computes that one op.

• builtin op: ops TFLite ships with (CONV_2D, ADD, SOFTMAX …), identified by an enum code.
• custom op: anything outside that set. It’s baked into the model as a string name, and to run it you have to register a kernel for that name at runtime.

Drop the attention model into a stock runtime and you get this:

Encountered unresolved custom op: Landmarks2TransformMatrix.
Node number 192 (Landmarks2TransformMatrix) failed to prepare.

A .tflite is a FlatBuffer, and each op in the graph points into an operator_codes table. Builtin ops are keyed by builtin_code (an enum number); custom ops are builtin_code = CUSTOM plus a custom_code string. So custom ops are matched by name. When the interpreter hits a custom_code it doesn’t recognize, it asks the resolver for a kernel of that name — and if there’s none, you get the unresolved custom op error above.

The names are right there in the file:

strings face_landmark_with_attention.tflite | grep -i transform
# TransformTensorBilinear
# TransformLandmarks
# Landmarks2TransformMatrix

The three custom ops this model uses

Op	What it does
Landmarks2TransformMatrix	Computes a transform matrix from landmark coordinates — e.g. using the coarse mesh’s points around the eye to derive that region’s crop/rotation.
TransformTensorBilinear	Bilinearly resamples the feature map per that matrix. Aligns and crops out the eye/lip regions.
TransformLandmarks	Applies the matrix to landmark coordinates. Maps the refined region coordinates back into the full-face frame.

Together they form the “attention” mechanism: get a coarse mesh → compute a per-region transform (Landmarks2TransformMatrix) → warp that region’s features (TransformTensorBilinear) → run the refinement head → map the result back to the original frame (TransformLandmarks). That’s what sharpens the eyes, lips, and irises.

Worth noting: these are not TFLite’s Flex (TF Select) ops. Flex pulls ops in from TF proper; these are separate kernels MediaPipe implements in its own source.

What a kernel looks like

In TFLite, one op kernel is a bundle of function pointers called TfLiteRegistration.

typedef struct {
  void* (*init)(TfLiteContext*, const char* buffer, size_t length);  // parse options
  void  (*free)(TfLiteContext*, void* buffer);
  TfLiteStatus (*prepare)(TfLiteContext*, TfLiteNode*);   // compute output shape, allocate
  TfLiteStatus (*invoke)(TfLiteContext*, TfLiteNode*);    // the actual compute
  const char* custom_name;                                // "TransformTensorBilinear"
  ...
} TfLiteRegistration;

A Register...V2() function returns this struct. So a custom op’s compute logic is the body of its prepare/invoke functions — and that has to be compiled into the binary.

If C function pointers feel foreign, the struct is clearer in Dart. Functions are first-class there, so you write a function-typed field instead of the (*prepare) notation. (Just an analogy — the real kernels are C++.)

typedef NodeFn = TfLiteStatus Function(TfLiteContext ctx, TfLiteNode node);

class TfLiteRegistration {
  NodeFn? prepare;  // not logic — an empty slot to plug an external function into
  NodeFn? invoke;
}

So it’s reg.prepare = myFn — you slot a function into an empty field. An op TFLite has never heard of runs the moment you slot its function in. That’s exactly how custom op registration works.

Why the stock runtime can’t run it

This package uses TFLite’s C API. When it creates an interpreter, the C API calls this internally:

// tensorflow/lite/create_op_resolver.h
std::unique_ptr<MutableOpResolver> CreateOpResolver();

The OpResolver it returns is the “op name/code → kernel” table. A stock build links create_op_resolver_with_builtin_ops, which hands back a resolver with builtin ops only.

So the distributed libtensorflowlite_c.so and TensorFlowLiteC.framework have neither the implementation nor the registration of those three custom ops.

strings libtensorflowlite_c.so | grep -c TransformTensorBilinear   # → 0

Which leaves building it myself.

The fix — swap out the resolver

Conveniently, the C API decides its entire op list through a single CreateOpResolver(). So replace that function with a version that registers “builtin + custom ops,” rebuild just the C binary, and the attention face mesh model loads.

MediaPipe already has the registration function. mediapipe/util/tflite/cpu_op_resolver.cc registers seven custom ops:

namespace mediapipe {
void MediaPipe_RegisterTfLiteOpResolver(tflite::MutableOpResolver* resolver) {
  resolver->AddCustom("MaxPoolingWithArgmax2D",
                      tflite_operations::RegisterMaxPoolingWithArgmax2D());
  resolver->AddCustom("MaxUnpooling2D",
                      tflite_operations::RegisterMaxUnpooling2D());
  resolver->AddCustom("Convolution2DTransposeBias",
                      tflite_operations::RegisterConvolution2DTransposeBias());
  resolver->AddCustom("TransformTensorBilinear",
                      tflite_operations::RegisterTransformTensorBilinearV2(),
                      /*version=*/2);
  resolver->AddCustom("TransformLandmarks",
                      tflite_operations::RegisterTransformLandmarksV2(),
                      /*version=*/2);
  resolver->AddCustom("Landmarks2TransformMatrix",
                      tflite_operations::RegisterLandmarksToTransformMatrixV2(),
                      /*version=*/2);
  resolver->AddCustom("Resampler", tflite_operations::RegisterResampler(),
                      /*version=*/1);
}
}  // namespace mediapipe

So the code we actually write is just this:

// mediapipe/custom_capi/create_op_resolver.cc
#include "mediapipe/util/tflite/cpu_op_resolver.h"
#include "tensorflow/lite/kernels/register.h"
#include "tensorflow/lite/mutable_op_resolver.h"

namespace tflite {

std::unique_ptr<MutableOpResolver> CreateOpResolver() {
  auto resolver = std::make_unique<ops::builtin::BuiltinOpResolver>();  // all builtins
  mediapipe::MediaPipe_RegisterTfLiteOpResolver(resolver.get());        // + custom ops
  return resolver;
}

}  // namespace tflite

This works because TFLite left an official extension point open. tensorflow/lite/c/BUILD has a c_api_without_op_resolver variant target that deliberately leaves CreateOpResolver() undefined:

tags = ["allow_undefined_symbols"],  # For tflite::CreateOpResolver().

Meaning “define it yourself and link it in.” Link the source above and you’re done.

Registration and implementation are different layers

Looking at that code, you might wonder: where’s the actual compute code for the op? We didn’t write it.

The key is that using an op takes two separate layers:

Layer	What it does	We write it?	Needed in binary?
Registration	name → kernel mapping (AddCustom(...))	Yes	Yes
Implementation	the prepare/invoke body (the actual math)	No (MediaPipe provides it)	Yes (mandatory)

AddCustom only wires a name to a function pointer. The actual code that pointer points to has to be compiled and linked somewhere; without it, the link blows up with undefined reference: RegisterTransformTensorBilinearV2.

Bazel pulls that implementation in through the deps chain. One thing first — this build runs inside the MediaPipe workspace. So in the chain below, //mediapipe/... is local source in the workspace we’re building, and @org_tensorflow//... is the TF dependency fetched from outside. (Why MediaPipe specifically is the next section.)

custom_create_op_resolver
  └─ //mediapipe/util/tflite:cpu_op_resolver              (registration function)
       └─ //mediapipe/util/tflite/operations:transform_tensor_bilinear
                                                          (TFLite kernel: prepare/invoke)
            └─ @org_tensorflow//tensorflow/lite/delegates/gpu/common/mediapipe:
               transform_tensor_bilinear                  (attribute structs, parsers, GPU kernel)

Had this been a genuinely new custom op rather than an existing MediaPipe one, I’d have had to write the prepare/invoke logic in C++ myself. Here MediaPipe already implemented it, so registration was all that was left.

Why the MediaPipe workspace, not stock TF

There’s a catch. That last line of the chain — tensorflow/lite/delegates/gpu/common/mediapipe/ — doesn’t exist in TF master.

These are files MediaPipe patches into TF at build time. Its WORKSPACE pins @org_tensorflow to a specific commit and applies four patches:

http_archive(
    name = "org_tensorflow",
    urls = ["https://github.com/tensorflow/tensorflow/archive/fad6b3cf5a7d51a437bd01ee929853bc8554b618.tar.gz"],
    patches = [
        "@//third_party:org_tensorflow_c_api_experimental.diff",
        "@//third_party:org_tensorflow_custom_ops.diff",     # ← this one injects the files
        "@//third_party:org_tensorflow_objc_build_fixes.diff",
        "@//third_party:org_tensorflow_verifier_int4.diff",
    ],
    ...
)

org_tensorflow_custom_ops.diff creates gpu/common/mediapipe/{BUILD, custom_parsers.cc, landmarks_to_transform_matrix.cc, transform_landmarks.cc, transform_tensor_bilinear.cc, ...} in the TF tree.

The upshot: the MediaPipe workspace is the one place where patched TF + op source + resolver all come together. That’s why the build runs inside the MediaPipe repo, with the C API target referenced as @org_tensorflow//....

One more note: @org_tensorflow isn’t the latest TF master but a commit pinned by MediaPipe’s WORKSPACE. The TF source doesn’t live in the MediaPipe repo either — it’s downloaded as a tar.gz during bazel build.

Here’s the whole picture:

[Source prep]
  Clone MediaPipe ── WORKSPACE fetches @org_tensorflow (pinned commit) + applies patches
                       └ custom_ops.diff → generates gpu/common/mediapipe/*
  + mediapipe/custom_capi/{create_op_resolver.cc, BUILD}   ← we add this

[Build (bazel, MediaPipe workspace)]
  custom_create_op_resolver ─┐
  c_api_without_op_resolver ─┼─► link ─► C API binary with custom ops
                             ┘
       Android: cc_binary(linkshared) → libtensorflowlite_c.so (per ABI)
       iOS:     ios_static_framework  → TensorFlowLiteC.framework

[Post-processing]
  strip / symbol hiding → verify → swap for the package's prebuilt binary

[Runtime]
  load model → TfLiteInterpreterCreate → CreateOpResolver() (builtin + custom)
             → AllocateTensors() succeeds

Android build

TF’s tensorflow/lite/tools/tflite-android.Dockerfile gives you a bazel + Android SDK + NDK environment. The image is just the environment; the source and custom ops are ours to add.

docker build . -t tflite-builder -f tflite-android.Dockerfile
docker run -d --name tflite-c \
  -v ~/attn-build/mediapipe:/mediapipe -v ~/attn-build:/host_dir \
  tflite-builder sleep infinity

docker exec tflite-c bash -lc '
cd /mediapipe
export ANDROID_HOME=/android/sdk ANDROID_NDK_HOME=/android/ndk
bazel build -c opt --config=android_arm64 --cxxopt=--std=c++17 \
    //mediapipe/custom_capi:libtensorflowlite_c.so
'

The errors I hit along the way:

1. The NDK toolchain isn’t registered in WORKSPACE. MediaPipe’s WORKSPACE loads android_ndk_repository but never calls it (# @unused). Leave it and there’s no Android toolchain, so the build won’t run. You have to add it:

android_ndk_repository(name = "androidndk", api_level = 30)
register_toolchains("@androidndk//:all")

2. The select() trap in the tflite_cc_shared_object macro. I first tried to use TF’s macro just like stock, but the bare //tensorflow:... labels inside it resolve, in the MediaPipe repo context, to @//tensorflow (= /mediapipe/tensorflow) and break with “no such package.”

The cause is Bazel’s label resolution. A bare // label with no repo resolves against the main repo running the build. The macro was written assuming its own workspace, where “the main repo is TF” — but we call it with mediapipe as the main repo (cd /mediapipe && bazel build), so //tensorflow:... lands at @mediapipe//tensorflow:..., which doesn’t exist. Worse, the macro’s select() for per-platform linkopts uses those same bare labels as its branch keys, so it’s not one dep that breaks but the whole branch.

So I dropped the macro and worked around it with a plain cc_binary(linkshared = 1). Pin every label to @org_tensorflow//... explicitly and there’s no bare //tensorflow and no select() left to misresolve. I restored the version-script the macro would have added by hand, via linkopts:

cc_binary(
    name = "libtensorflowlite_c.so",
    linkshared = 1,
    linkopts = [
        "-Wl,--version-script,$(location @org_tensorflow//tensorflow/lite/c:version_script.lds)",
        "-Wl,-z,max-page-size=16384",
    ],
    deps = [
        ":custom_create_op_resolver",
        "@org_tensorflow//tensorflow/lite/c:c_api_without_op_resolver",
        "@org_tensorflow//tensorflow/lite/c:c_api_experimental_without_op_resolver",
        "@org_tensorflow//tensorflow/lite/c:version_script.lds",
    ],
)

3. XNNPACK’s newest ISA kernels don’t build on x86_64. NDK r25b’s clang doesn’t know flags like -mavx512fp16 or -mavxvnniint8. Just disable those kernels:

--define=xnn_enable_avx512fp16=false --define=xnn_enable_avxvnniint8=false ...

4. The output comes out 2.5× the size of stock. My first build was 10MB; stock is 4MB. “Did the custom ops bloat it?” No — stock was simply already stripped.

llvm-strip --strip-all libtensorflowlite_c.so
# arm64-v8a: 10MB → 4.1MB, x86_64: 11.5MB → 6.1MB (same as stock)

5. 16KB page alignment. The official TFLite release .so is linked with -Wl,-z,max-page-size=16384, so its LOAD segments are 16KB (0x4000) aligned. My bazel build didn’t have that linkopt, so they came out 4KB (0x1000) aligned.

For context, Google Play requires 16KB page support on 64-bit devices for new apps and updates targeting Android 15+ as of November 1, 2025. When you’re swapping a prebuilt native library, that’s not something you can wave off.

readelf -lW libtensorflowlite_c.so | grep LOAD   # Align must be 0x4000

Verification

Three things to check on the build:

readelf -lW $f | grep LOAD                              # 0x4000 alignment
strings -n6 $f | grep -c TransformTensorBilinear        # custom op present
nm -D --defined-only $f | grep -c TfLiteXNNPackDelegate  # XNNPACK symbols preserved

A name showing up in strings doesn’t mean the op actually runs. An arm64 .so won’t run on x86 Linux, so I also built the same custom C API for host and ran a small C program that took the model through TfLiteInterpreterCreate → AllocateTensors. It cleared the exact point a stock build would fail at, and produced seven output tensors, all correct:

Output tensor	Size	Meaning
mesh	1404	468 landmarks × 3 (x, y, z)
lips	160	80 × 2
left eye / right eye	142 each	71 × 2
left iris / right iris	10 each	5 × 2
face flag	1	face-presence probability

So it’s confirmed not just that “the op is in there” but that “it runs.”

iOS build

The idea is identical to Android. The resolver source is platform-independent, so it’s reused as-is. The differences: the output is a TensorFlowLiteC.framework instead of a .so, and you need macOS + Xcode.

The plan: mirror the stock framework target → failed

I tried what worked on Android: mirror the stock //tensorflow/lite/ios:TensorFlowLiteC_framework (the tflite_ios_framework macro) and just swap the resolver in its deps. Two walls stopped me.

1. BUILD.apple isn’t a bazel package. TF’s tensorflow/lite/ios/ names its build file BUILD.apple, not BUILD, so bazel doesn’t recognize the directory as a package. load(".../tensorflow/lite/ios:ios.bzl"), the header targets, allowlist_TensorFlowLiteC.txt — all fail with “not a package.”

2. Visibility + a dependency cycle. That package’s header targets and tools (extract_object_files, hide_symbols_with_allowlist) are visible only to //tensorflow/lite:__subpackages__, so @mediapipe can’t use them. And I can’t move my target into @org_tensorflow either — my resolver depends on //mediapipe/util/tflite:cpu_op_resolver, which would make a TF → MediaPipe cycle. “Just patch TF’s BUILD” was dead for the same cyclic reason.

The fix: assemble the framework yourself

Since I couldn’t use TF’s iOS macro, I assembled it self-contained inside @mediapipe with rules_apple’s public ios_static_framework. Because my resolver depends on //mediapipe/util/tflite:cpu_op_resolver, the framework target has to live in @mediapipe too — and there it references none of TF’s ios package, linking everything from local cc source (the C API shell + resolver). Same idea as the plain cc_binary workaround on Android.

In fact both platforms share the same skeleton: the same MediaPipe clone, our own BUILD target inside it, assembled with a generic bazel rule. Only the rule and the output differ.

	Android	iOS
Build workspace	@mediapipe (clone)	@mediapipe (same clone)
Our target location	mediapipe/custom_capi/BUILD	mediapipe/custom_capi/BUILD
Rule used	built-in cc_binary	rules_apple ios_static_framework
TF macro	none (worked around)	none (worked around)
Output	libtensorflowlite_c.so	TensorFlowLiteC.framework

One thing worth knowing: an iOS .framework isn’t a single file but a bundle (directory) holding a binary + Headers/ + Modules/. A custom op only changes the code (the binary), not the C API signatures — so the only thing that needs swapping is the Mach-O binary inside, and stock’s Headers/Modules stay valid. But the ios_static_framework rule requires hdrs. So I copied the existing framework’s Headers/ in just to satisfy the rule and get it to build, then threw away the produced bundle’s headers/modules, pulled out only the binary, and reused the existing bundle’s Headers/Modules. The swap itself is a one-file change — just the TensorFlowLiteC binary inside the framework folder.

JAVA_HOME=/opt/homebrew/opt/openjdk USE_BAZEL_VERSION=7.4.1 bazel build \
    --config=ios --ios_multi_cpus=arm64,x86_64 \
    -c opt --cxxopt=--std=c++17 \
    --repo_env=HERMETIC_PYTHON_VERSION=3.12 \
    --repo_env=JAVA_HOME=/opt/homebrew/opt/openjdk \
    //mediapipe/custom_capi:TensorFlowLiteC_framework

Every flag has a reason:

• --config=ios_fat (the documented one) breaks. ios_fat includes armv7/i386, and the bzlmod rules_apple I got has no ios_armv7 mapping, so it throws a transition error (key "ios_armv7" not found). I used --config=ios --ios_multi_cpus=arm64,x86_64 instead.
• HERMETIC_PYTHON_VERSION=3.12 — system python 3.14 is outside TF’s requirements_lock range (3.9–3.12) and fails.
• JAVA_HOME + --repo_env=JAVA_HOME — with no JDK, rules_java’s local_jdk repo breaks (no such package @rules_java~//tools/jdk) and the apple bundling analysis dies.

367MB → 18MB: symbol hiding was mandatory

The binary ios_static_framework spat out was 367MB — against stock’s 47MB. The reason: it’s just an ar archive with every symbol global. Stock’s tflite_ios_framework macro runs a symbol merge/hide step internally; without the macro, that step was missing entirely.

So I ran stock’s script (hide_symbols_with_allowlist.sh) by hand. What it does:

• per arch, ld -r -exported_symbols_list <allowlist> -x to merge all objects into a single Mach-O, and
• localize/strip every symbol outside the allowlist (_TfLite*, *AcquireFlexDelegate*).

The result: 18MB, exactly 219 exported _TfLite* symbols, zero non-TfLite externals. So it wasn’t “custom ops made it bigger” — it was the macro’s post-processing being absent that made it look bigger. The iOS version of what strip did in one shot on Android, only far more involved.

Why it’s even smaller than stock (47MB)

The stock framework was 47.5MB before the swap; mine is 18.1MB. Both are arm64+x86_64, so it’s not slice count. Two things seem to overlap:

1. A narrower link than stock. My deps are just c_api_without_op_resolver + c_api_experimental_without_op_resolver + resolver, while stock’s :tensorflow_lite_c piles on profiling/telemetry and more.
2. Fewer symbols and metadata. Stock carries the SUBSECTIONS_VIA_SYMBOLS flag with ~2,520 symbol strings; mine has no such flag and ~646. Merging every object into one Mach-O with ld -r … -x and stripping local symbols shrinks the symbol and relocation tables a lot. It’s less that code shrank and more that metadata did.

Verification — “the Runner binary has no symbols”

I put the framework in ios/Frameworks/, built the example, and checked symbols — and the Runner binary showed zero TfLite symbols, so I thought it had failed.

The cause was how this package loads. On Apple, src/tflite_runtime.h finds symbols via dlsym(RTLD_DEFAULT, ...) — a static-link model — so the TfLite symbols end up not in the Runner but in the plugin’s dynamic framework (Runner.app/Frameworks/mediapipe_face_mesh.framework). Looking there, TfLiteInterpreterCreate and the custom ops were present and fine.

Android vs iOS in one line: Android was nothing more than “dodge the macro’s select trap with a plain cc_binary,” while iOS meant I couldn’t reach TF’s iOS package at all (BUILD.apple + visibility/cycle), so I had to assemble the framework myself and even restore, by hand, the symbol hiding the macro would have done.

On the package side

The rest is ordinary integration:

• One option added (enableAttentionMesh, off by default, so existing behavior is unchanged).
• When it’s on, the model’s seven output tensors are assembled into the 478 layout: lay down the 468 mesh, overwrite the eye/lip indices, and append the 10 iris points. The assembly rules follow MediaPipe’s tensors_to_face_landmarks_with_attention.pbtxt config.
• Everything downstream — coordinate transforms, ROI, blendshapes — can’t tell the difference, since it just receives 478 points.

Wrapping up

• A .tflite is an op graph, and a custom op is one outside the standard set that you have to register a kernel for at runtime.
• The attention face mesh is built from three MediaPipe custom ops (transform-matrix computation / feature warp / landmark inverse-transform).
• The C API decides its op list through CreateOpResolver(), and stock has builtins only. → Replace just that function and rebuild the C binary, and the model loads with no app-code changes. It’s the official extension point TFLite left open via c_api_without_op_resolver.
• That rebuild has to happen inside the MediaPipe workspace (patched, pinned TF), because the custom ops’ dependency files don’t exist in TF master — only through MediaPipe’s patch.
• The custom ops themselves are plain TFLite kernels, so porting was clean. What actually ate the time was the build system, not the ops.
• Size is easy to misread. Custom ops don’t bloat anything; it only looked bigger because we skipped the post-processing (strip / symbol hiding) stock does for you. Restore it and you’re at Android 4.1MB (= stock), iOS 18MB (< stock’s 47MB).

References

• MediaPipe cpu_op_resolver.cc
• MediaPipe WORKSPACE (pinned TF commit + patches)
• TFLite tensorflow/lite/c/BUILD (c_api_without_op_resolver)
• TFLite tensorflow/lite/ios/ (BUILD.apple, hide_symbols_with_allowlist.sh)
• Support 16 KB page sizes - Android Developers
• mediapipe_face_mesh (pub.dev)