jax to tflite conversion

[ASEM000]

The following fails if I specify the backend as cpu/gpu with the following error:

'stablehlo.custom_call' op is not part of the stablehlo support yet.

import jax
import jax.numpy as jnp
import tensorflow as tf

if __name__ == "__main__":
    x_input = jnp.zeros((4624, 3468, 3))
    predict = jax.jit(lambda x: x, backend='cpu')
    converter = tf.lite.TFLiteConverter.experimental_from_jax(
    [predict], [[('input1', x_input)]])
    tflite_model = converter.convert()
    # check
    interpreter = tf.lite.Interpreter(model_content=tflite_model)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    interpreter.set_tensor(input_details[0]['index'], x_input)
    interpreter.invoke()
    result = interpreter.get_tensor(output_details[0]['index'])

What is the recommended way to convert jitted functions to tflite?

[jakevdp]

@gnecula might have suggestions here.

[ASEM000]

@gnecula can you provide some pointers?

[gnecula]

JAX on cpu/gpu does use custom calls occasionally. I think that for the core JAX primitives the TFLite converter should handle them, so this is a request that should be made to the TFLite team.

[ASEM000]

Will do, thanks