Frontend Integration

def compile_saved_model(saved_model_dir: str, **kwargs):
    """
    Compiles an on-disk saved model to an IREE binary.
    
    Parameters:
    - saved_model_dir: Path to directory where the model was saved
    - **kwargs: Keyword args corresponding to ImportOptions or CompilerOptions
    
    Returns:
    bytes | None: Compiled output or None if output_file was specified
    """

def compile_module(module, saved_model_dir: Optional[str] = None, **kwargs):
    """
    Compiles a tf.Module to an IREE binary (by saving to disk).
    
    Parameters:
    - module: The tf.Module instance to convert to MLIR
    - saved_model_dir: Optional path to save the tf.Module to
    - **kwargs: Keyword args corresponding to ImportOptions or CompilerOptions
    
    Returns:
    bytes | None: Same as compile_saved_model()
    """

def is_available() -> bool:
    """
    Determine if TensorFlow and the TF frontend are available.
    
    Returns:
    bool: True if TensorFlow integration is available
    """

class ImportType(Enum):
    """
    Import type of the model.
    """
    OBJECT_GRAPH = "savedmodel_v2"
    V2 = "savedmodel_v2"
    SIGNATURE_DEF = "savedmodel_v1"
    V1 = "savedmodel_v1"
    
    @staticmethod
    def parse(spec: Union[str, ImportType]) -> ImportType:
        """
        Parses or returns an ImportType.
        
        Parameters:
        - spec: ImportType instance or case-insensitive name
        
        Returns:
        ImportType: Parsed import type
        """

class ImportOptions(CompilerOptions):
    """
    Import options layer on top of the backend compiler options.
    
    Attributes:
    - exported_names: Sequence[str], exported names to keep (object graph/v2 models only)
    - import_only: bool, only import without compilation (default: False)
    - import_type: ImportType, type of import to perform (default: OBJECT_GRAPH)
    - input_type: Union[InputType, str], input type (default: STABLEHLO_XLA)
    - saved_model_tags: Set[str], tags to export (signature def/v1 models only)
    - save_temp_iree_input: Optional[str], save IR result of import
    """
    
    def __init__(self, **kwargs): ...

def compile_file(fb_path: str, **kwargs):
    """
    Compiles a TFLite FlatBuffer file to an IREE binary.
    
    Parameters:
    - fb_path: Path to the FlatBuffer
    - **kwargs: Keyword args corresponding to ImportOptions or CompilerOptions
    
    Returns:
    bytes | None: Compiled output or None if output_file was specified
    """

def compile_str(fb_content: bytes, **kwargs):
    """
    Compiles TFLite FlatBuffer content to an IREE binary.
    
    Parameters:
    - fb_content: FlatBuffer content as bytes
    - **kwargs: Keyword args corresponding to ImportOptions or CompilerOptions
    
    Returns:
    bytes | None: Compiled output or None if output_file was specified
    """

def is_available() -> bool:
    """
    Determine if TensorFlow and the TFLite frontend are available.
    
    Returns:
    bool: True if TFLite integration is available
    """

class ImportOptions(CompilerOptions):
    """
    Import options layer on top of the backend compiler options.
    
    Attributes:
    - input_arrays: Sequence[str], input array node names (if different from default)
    - output_arrays: Sequence[str], output array node names (if different from default)
    - import_only: bool, only import without compilation (default: False)
    - import_extra_args: Sequence[str], extra arguments to pass to import tool
    - save_temp_tfl_input: Optional[str], save IR from importing flatbuffer
    - save_temp_iree_input: Optional[str], save IR result of import
    - input_type: Optional[str], input type (default: "tosa")
    """
    
    def __init__(self, **kwargs): ...

DEFAULT_TESTING_BACKENDS: List[str] = ["llvm-cpu"]

import iree.compiler.tools.tf as tf_compiler

# Basic SavedModel compilation
bytecode = tf_compiler.compile_saved_model(
    "path/to/saved_model",
    target_backends=["llvm-cpu"]
)

# Advanced compilation with export selection
bytecode = tf_compiler.compile_saved_model(
    "path/to/saved_model",
    exported_names=["serving_default", "predict"],
    target_backends=["llvm-cpu", "cuda"],
    import_type=tf_compiler.ImportType.OBJECT_GRAPH,
    optimize=True
)

# Import-only workflow
mlir_content = tf_compiler.compile_saved_model(
    "path/to/saved_model",
    import_only=True,
    save_temp_iree_input="/tmp/model.mlir"
)

import tensorflow as tf
import iree.compiler.tools.tf as tf_compiler

# Create a TensorFlow module
class SimpleModule(tf.Module):
    def __init__(self):
        super().__init__()
        self.dense = tf.keras.layers.Dense(10)
    
    @tf.function(input_signature=[tf.TensorSpec([None, 784], tf.float32)])
    def __call__(self, x):
        return self.dense(x)

# Compile the module
module = SimpleModule()
bytecode = tf_compiler.compile_module(
    module,
    target_backends=["llvm-cpu"],
    export_names=["__call__"]
)

# Save module to disk and compile
tf_compiler.compile_module(
    module,
    saved_model_dir="/tmp/my_model",
    target_backends=["cuda"],
    output_file="model.vmfb"
)

import iree.compiler.tools.tflite as tflite_compiler

# Basic TFLite compilation
bytecode = tflite_compiler.compile_file(
    "model.tflite",
    target_backends=["llvm-cpu"]
)

# Advanced compilation with array specification
bytecode = tflite_compiler.compile_file(
    "model.tflite",
    input_arrays=["input_tensor"],
    output_arrays=["output_tensor"],
    target_backends=["vulkan-spirv"],
    optimize=True,
    save_temp_iree_input="/tmp/tflite_imported.mlir"
)

# Import-only workflow
mlir_content = tflite_compiler.compile_file(
    "model.tflite",
    import_only=True
)

import iree.compiler.tools.tf as tf_compiler
import iree.compiler.tools.tflite as tflite_compiler

# Check TensorFlow availability
if tf_compiler.is_available():
    print("TensorFlow integration available")
    try:
        bytecode = tf_compiler.compile_saved_model(
            "path/to/model",
            target_backends=["llvm-cpu"]
        )
    except Exception as e:
        print(f"TensorFlow compilation failed: {e}")
else:
    print("TensorFlow integration not available")

# Check TensorFlow Lite availability
if tflite_compiler.is_available():
    print("TensorFlow Lite integration available")
    try:
        bytecode = tflite_compiler.compile_file(
            "model.tflite",
            target_backends=["llvm-cpu"]
        )
    except Exception as e:
        print(f"TFLite compilation failed: {e}")
else:
    print("TensorFlow Lite integration not available")

import iree.compiler.tools.tf as tf_compiler

# Custom import options for TensorFlow
import_options = tf_compiler.ImportOptions(
    exported_names=["serving_default"],
    import_type=tf_compiler.ImportType.V2,
    target_backends=["llvm-cpu", "cuda"],
    optimize=True,
    output_mlir_debuginfo=True,
    save_temp_iree_input="/tmp/tf_imported.mlir"
)

# Compile with custom options
bytecode = tf_compiler.compile_saved_model(
    "path/to/saved_model",
    **import_options.__dict__
)

import iree.compiler.tools.tf as tf_compiler
import os

def compile_for_multiple_targets(saved_model_path: str, targets: list):
    """Compile a SavedModel for multiple target backends."""
    results = {}
    
    for target in targets:
        print(f"Compiling for {target}...")
        
        try:
            output_file = f"model_{target}.vmfb"
            tf_compiler.compile_saved_model(
                saved_model_path,
                target_backends=[target],
                output_file=output_file,
                optimize=True
            )
            results[target] = output_file
            print(f"Successfully compiled for {target}: {output_file}")
            
        except Exception as e:
            print(f"Failed to compile for {target}: {e}")
            results[target] = None
    
    return results

# Usage
targets = ["llvm-cpu", "cuda", "vulkan-spirv", "rocm"]
results = compile_for_multiple_targets("path/to/saved_model", targets)