HLO protobuf to pytorch / tensorflow graph

Assume we have HLO protobuf from a model through Pytorch-XLA or Tensorflow.

1. Is there a way to create computational graph from it?
2. Is it possible to create Pytorch-XLA and Tensorflow model from it?

In python, input has following type

hlo_pb2.HloModuleProto()

I think one approach would be to follow this two0step process.

1. Extract relevant information from this protobuf object, such as the nodes, their connections, and attributes, and proceed with the conversion process into ONNX
2. TensorFlow provides built-in utilities for converting ONNX models into TensorFlow models. After converting the HLO protobuf to ONNX format using the appropriate tools or libraries. Then, use TensorFlow's tf.compat.v1.graph_util.import_graph_def() or tf.saved_model.loader.load() to load the ONNX file and obtain a TensorFlow computational graph.

Here is a sample snippet that you would need to adjust to achieve step 1

Here's an example code snippet that demonstrates how to extract information from an HLO protobuf object and convert it into an ONNX file using PyTorch:

import torch
import onnx
from onnx import helper
from onnx import AttributeProto, TensorProto, GraphProto

def convert_hlo_to_onnx(hlo_module):
    # Create an empty ONNX graph
    graph = helper.make_graph([], "hlo_model", [], [])

    # Keep track of the ONNX node names and their corresponding outputs
    node_outputs = {}

    # Iterate over the HLO computations and instructions
    for computation in hlo_module.computations:
        # Iterate over the HLO instructions in the computation
        for instruction in computation.instructions:
            # Extract instruction attributes
            instruction_name = instruction.name
            instruction_opcode = instruction.opcode
            instruction_outputs = instruction.operand

            # Create ONNX node with corresponding inputs and outputs
            onnx_node = helper.make_node(instruction_opcode, instruction.operand, instruction_outputs, name=instruction_name)

            # Set any additional attributes for the ONNX node if needed
            # For example, if the HLO instruction has attributes 'attr1' and 'attr2',
            # you can add them to the ONNX node as follows:
            # onnx_node.attribute.extend([
            #     helper.make_attribute('attr1', instruction.attr1),
            #     helper.make_attribute('attr2', instruction.attr2)
            # ])

            # Add the ONNX node to the graph
            graph.node.extend([onnx_node])

            # Update the node_outputs dictionary with the current ONNX node outputs
            node_outputs[instruction_name] = instruction_outputs

    # Iterate over the HLO computations again to establish connections between nodes
    for computation in hlo_module.computations:
        # Iterate over the HLO instructions in the computation
        for instruction in computation.instructions:
            # Get the current instruction name and its outputs
            instruction_name = instruction.name
            instruction_outputs = instruction.operand

            # Iterate over the outputs and connect them to the corresponding inputs
            for output in instruction_outputs:
                # Check if the output is used as an input in any subsequent instruction
                if output in node_outputs:
                    # Get the corresponding ONNX node output name
                    onnx_node_output = node_outputs[output]

                    # Find the ONNX node with the same name as the current instruction
                    onnx_node = next(node for node in graph.node if node.name == instruction_name)

                    # Update the inputs of the ONNX node to connect with the output of the previous node
                    onnx_node.input.remove(output)
                    onnx_node.input.extend(onnx_node_output)

    # Create the ONNX model with the graph
    model = helper.make_model(graph)

    # Save the ONNX model to a file
    onnx.save_model(model, "converted_model.onnx")

# Assuming you have an HLO module protobuf object named 'hlo_module'
hlo_module = hlo_pb2.HloModuleProto()

# Call the conversion function
convert_hlo_to_onnx(hlo_module)

In this example, we ie iterate over the computations and instructions fields of the HloModuleProto. Each instruction represents an HLO instruction within a computation. We extract the instruction name, opcode, and outputs. Then, we create an ONNX node using helper.make_node() and set the inputs, outputs, and attributes accordingly. We add the ONNX node to the graph and keep track of the node outputs in the node_outputs dictionary.

After creating the ONNX graph, we use helper.make_model() to create an ONNX model with the graph, and finally, onnx.save_model() is used to save the ONNX model to a file ("converted_model.onnx" in this example).

Please note that this code assumes you have the necessary dependencies installed, including the onnx and torch packages. Also, make sure to import the relevant modules (onnx, helper, AttributeProto, TensorProto, GraphProto) from the onnx package.

Remember to adapt this code to fit your specific HLO protobuf structure and attributes.