OPTIMIZED_HLO_SERIALIZED,
  OPTIMIZED_HLO_PROTO_SERIALIZED,
  OPTIMIZED_HLO_DOT
};

struct ArgShapeAndDType {
  TensorShape shape;
  DataType dtype;
};

enum class CompilerArgSource {
  TENSOR_SPEC,
  CONCRETE_INPUT,
};

// Returns the IR format of the selected stage for a given function `func_name`
// using library runtime `runtime` on a device `dev` with given
// `inputs_arg_shape_and_dtype` and `input_handles`.

    const mlir::Type& type) {
  DataType dtype;
  TF_RETURN_IF_ERROR(ConvertToDataType(type, &dtype));
  TensorSpecProto tensor_spec;
  tensor_spec.set_dtype(dtype);
  *tensor_spec.mutable_shape() = ConvertTypeToTensorShape(type).AsProto();
  return tensor_spec;
}

// Converts the tensor shape proto into an MLIR shape attribute.
absl::StatusOr<mlir::Attribute> ConvertTensorShapeProto(

  # CHECK-SAME: ) -> (
  # CHECK-SAME:   tensor<1xf32> {tf_saved_model.index_path = ["x"]})
  # CHECK-SAME: attributes {{.*}} tf_saved_model.exported_names = ["f0006_multiple_return_statements"]
  @tf.function(input_signature=[tf.TensorSpec([], tf.float32)])
  def f0006_multiple_return_statements(self, x):
    if x > 3.:
      return {'x': tf.constant(1.0, shape=[1])}
    else:
      return {'x': tf.constant(1.0, shape=[1])}

    TF_RETURN_IF_ERROR(LockVariables(absl::MakeSpan(variable_infos)));
  }

  absl::StatusOr<std::vector<XlaCompiler::Argument>> args;

  if (compiler_arg_source == CompilerArgSource::TENSOR_SPEC) {
    args = BuildXlaCompilerArgumentFromTensorSpec(fbody, constant_arg_indices,
                                                  inputs, variable_infos,

    model = TwoMatmulModel()
    input_shape = (1, 2)

    save.save(
        model,
        self._input_saved_model_path,
        signatures=model.matmul.get_concrete_function(
            tensor_spec.TensorSpec(
                shape=input_shape, dtype=dtypes.float32, name='input_tensor'
            )
        ),
    )

    def data_gen() -> repr_dataset.RepresentativeDataset:
      for _ in range(50):

#include "tensorflow/c/experimental/saved_model/public/signature_def_param.h"
#include "tensorflow/c/experimental/saved_model/public/signature_def_param_list.h"
#include "tensorflow/c/experimental/saved_model/public/tensor_spec.h"
#include "tensorflow/c/tf_datatype.h"
#include "tensorflow/c/tf_shape.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_tensor.h"
#include "tensorflow/core/lib/io/path.h"

    *   Example usage to create server: ```python server =
        tf.distribute.experimental.rpc.Server.create("grpc", "127.0.0.1:1234")
        @tf.function(input_signature=[ tf.TensorSpec([], tf.int32),
        tf.TensorSpec([], dtypes.int32) ]) def _remote_multiply(a, b): return
        tf.math.multiply(a, b)

        server.register("multiply", _remote_multiply) ```

}
'''

@tf.function
def foo(x, y):
  return mlir_passthrough_op([x, y], mlir_module, Toutputs=[tf.float32])

graph_def = foo.get_concrete_function(tf.TensorSpec([10], tf.float32), tf.TensorSpec([10], tf.float32)).graph.as_graph_def()
```
  }];

  let arguments = (ins
    Variadic<TF_Tensor>:$inputs,

    StrAttr:$mlir_module
  );