jax / jaxlib

[Litsardos]

Could you add this package in the https://pypi.flet.dev/?

[ndonkoHenri]

Thanks for the request! Unfortunately jax/jaxlib can't be added to pypi.flet.dev — this one is blocked well below the packaging layer, so let me explain why, and then show the path that does work, because jax models can absolutely still run in a Flet app.

Why a recipe isn't possible

• There's no source to build from: jaxlib (the native engine jax requires) has never published a source distribution — it ships wheels only, built exclusively with Bazel inside Google's XLA monorepo, and that build system has no Android or iOS support at all. Everything on pypi.flet.dev is cross-compiled from source; jax is in the same category as tensorflow/torch here.
• iOS couldn't run it even if a wheel existed: XLA's CPU backend executes your jax code by JIT-compiling machine code at runtime, and iOS forbids third-party apps from doing that (apps can't create writable+executable memory). There's no interpreter fallback in jaxlib — jax.disable_jit() still runs XLA-compiled code — so this isn't something packaging work can fix.
• Size: jaxlib unpacks to ~330 MB per CPU architecture (one 320 MB shared library containing XLA + LLVM). That's far beyond a sensible mobile app payload.
• And import jax hard-requires jaxlib at the exact pinned version — there's no reduced or pure-Python mode.

What works instead: export the model, run it with a mobile runtime

Google's own supported answer for "jax on mobile" is to keep jax on your desktop for training/tracing, export the function, and run the exported model on-device with a lightweight runtime. Both runtimes you'd need are already on pypi.flet.dev for iOS and Android.

Option A — TFLite (the official jax path, guide)

On your desktop:

import tensorflow as tf
from jax.experimental import jax2tf

tf_fn = tf.function(
    jax2tf.convert(lambda x: predict(trained_params, x), with_gradient=False),
    input_signature=[tf.TensorSpec([1, 28, 28], tf.float32)],
    autograph=False,
)
converter = tf.lite.TFLiteConverter.from_concrete_functions([tf_fn.get_concrete_function()])
open("model.tflite", "wb").write(converter.convert())

Then in your Flet app (dependencies = ["flet", "tflite-runtime", "numpy"]), with model.tflite bundled next to your app code (guide):

import os
import numpy as np
from tflite_runtime.interpreter import Interpreter

interpreter = Interpreter(os.path.join(os.path.dirname(__file__), "model.tflite"))
interpreter.allocate_tensors()
inp, out = interpreter.get_input_details()[0], interpreter.get_output_details()[0]
interpreter.set_tensor(inp["index"], x.astype(np.float32))
interpreter.invoke()
result = interpreter.get_tensor(out["index"])

Option B — ONNX

Export with jax2onnx and run with onnxruntime from pypi.flet.dev (also both platforms). Same idea, different format — pick whichever converter covers your model.

One honest caveat: both converters cover a large-but-not-total subset of jax ops, so validate the converted model on desktop (compare outputs against jax) before wiring it into the app.

If you need jax-like APIs on the device itself

• keras 3 with the numpy backend is on pypi.flet.dev — layers/inference with a jax-adjacent API, validated on Android (note: the numpy backend has no CNN support).
• For autodiff at runtime, autograd (jax's pure-Python predecessor — numpy-based grad) installs straight from PyPI with no recipe needed, and numpy is already on pypi.flet.dev.

If you share what you're using jax for (model inference? optimization? autodiff?), I can point you at the closest supported combination.