XlaCallModule or another way of loading StableHLO code about coremltools HOT 4 OPEN

Hey @polvalente, could you please elaborate your use case a little bit?

1. How do you get your StableHLO module in the first place? (By training a JAX model I guess?)
2. What does the TF graph look like? Does it fully digest the StableHLO module, or is XlaCallModule a high level op that still delegates to the StableHLO graph?
3. Could you please share the original error message?
4. Could you please provide a minimal reproduce?

from coremltools.

@YifanShenSZ Thanks for the response!

I'll preface this by saying that my end goal is to have some way to compile https://github.com/elixir-nx/nx into CoreML for running on iOS devices. I tried StableHLO because that seemed to be a good/ready pathway.

1. It's close to how you would get a Jax model, but I'm actually using https://github.com/elixir-nx/nx/tree/main/exla, which uses openxla to compile Elixir (through Nx) code into a StableHLO module.
2. XlaCallModule is a high level op that contains/delegates to the StableHLO graph. I unfortunately wasn't able to find a way to fully digest a stablehlo module into a TF graph.
   3 and 4. I'll get back to you shortly with a reproduction script and sample output.

from coremltools.

I see, since XlaCallModule is a delegation, unfortunately we won't be able to convert it as is.

However, I do think if you can let TF graph digest the StableHLO module, then we may have a chance to convert. One potential way to do so is to mimic JAX2TF, although you may need to dig into how MLIR is called under the hood

from coremltools.

@YifanShenSZ Unfortunately I think JAX2TF does exactly what I'm doing. Thanks for the suggestion though!

Here's the reproduction script and sample output files anyway.

import tensorflow as tf
from tensorflow.compiler.tf2xla.python import xla as tfxla
import coremltools as ct


stablehlo_module = """
module {
  func.func public @main(%arg0: tensor<2xi64>, %arg1: tensor<2xi64>) -> tensor<2xi64> {
    %0 = stablehlo.create_token : !stablehlo.token
    %1 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<2xi64>) -> tensor<2xi64>
    %2 = stablehlo.broadcast_in_dim %arg1, dims = [0] : (tensor<2xi64>) -> tensor<2xi64>
    %3 = stablehlo.multiply %1, %2 : tensor<2xi64>
    %4 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<2xi64>) -> tensor<2xi64>
    %5 = stablehlo.broadcast_in_dim %arg1, dims = [0] : (tensor<2xi64>) -> tensor<2xi64>
    %6 = stablehlo.subtract %4, %5 : tensor<2xi64>
    %7 = stablehlo.broadcast_in_dim %3, dims = [0] : (tensor<2xi64>) -> tensor<2xi64>
    %8 = stablehlo.broadcast_in_dim %6, dims = [0] : (tensor<2xi64>) -> tensor<2xi64>
    %9 = stablehlo.add %7, %8 : tensor<2xi64>
    stablehlo.return %9 : tensor<2xi64>
  }
}
"""

# Define the computation using tf.function
@tf.function(input_signature=[
    tf.TensorSpec(shape=[2], dtype=tf.int64, name='arg0'),
    tf.TensorSpec(shape=[2], dtype=tf.int64, name='arg1')
])
def model_fn(arg0, arg1):
    args = [arg0, arg1]
    res = tfxla.call_module(
        args,
        module=stablehlo_module,
        Tout=[tf.int64],
        Sout=[(2,)],
        function_list=[],
        platforms=[],
    )
    return res

# Convert the SavedModel to Core ML
coreml_model = ct.convert(model_fn.get_concrete_function().graph.as_graph_def(), source='tensorflow')

sample_output.txt

I am able to run the model_fn function as if it were a standard TF model, so there's that at least :)

Feel free to close this is you do feel there's no path forward here!

from coremltools.