# fastapi_jax_server.py from fastapi import FastAPI, File, UploadFile import jax.numpy as jnp from jax import jit from PIL import Image import numpy as np app = FastAPI() Assume we have a function predict_jax @jit def preprocess_image(image_array): # Normalize, resize, etc. return jnp.array(image_array)
@app.post("/predict") async def predict(file: UploadFile = File(...)): img = Image.open(file.file).resize((224, 224)) img_array = np.array(img) / 255.0 jax_input = preprocess_image(img_array) output = predict_jax(jax_input) # jax array return "prediction": float(output) # convert to Python float