How to use jax2tf with ffi?

[krzysztofrusek]

I tried to used jax2tf on a function using jnp.linalg.svd and it cannot be used in tf due to missing ffi call to lapack.

Is there a way around? Maybe disabling lowering to ffi call in svd?

Code example

import jax
import jax.numpy as jnp
from jax.experimental import jax2tf

import tensorflow as tf

model = tf.Module()
model.f = tf.function(jax2tf.convert(lambda x: jnp.linalg.svd(x,False)), autograph=False,
                             input_signature=[tf.TensorSpec(shape=[4,14], dtype=tf.float32)])


output_dir = 'pca'
tf.saved_model.save(model, output_dir,options=tf.saved_model.SaveOptions(experimental_custom_gradients=False))
restored_model = tf.saved_model.load(output_dir)
x = tf.zeros((4,14))
print(restored_model.f(x))

when I change jnp.linalg.svd to e.g. jnp.sum this works.
I believe there is no issue on GPU as svd is lowered to xla op.
I am looking for a simple solution like getting pre-ffi implementation of svd instead of lapac call.

FYI 1. the code is going to be executed in java with TensorFLow binding.
FYI2. We planed to implement all numerical algorithms in jax and use saved model to bring them to application so there may be another opps that get lowered to ffi call.