yobx.tensorflow.ops.xla_call_module_parsing

StableHLO MLIR → layer-dict parser for the XlaCallModule converter.

parse_mlir() is the public entry point. Low-level extraction utilities that have no dependency on parsing state live in xla_call_module_helper.

yobx.tensorflow.ops.xla_call_module_parsing.parse_mlir(mlir_string: str) → List[dict]

Parse a StableHLO MLIR module text into a list of layer dicts.

Each dict has at minimum the keys id, op, operands, shape, and loc. Recognised op values include:

• "Input" – a function argument (tensor input).

• "return" – function return.

• "constant" – stablehlo.constant; has extra key dense_content.

• "dot_general" – matrix multiply; has lhs_contracting / rhs_contracting.

• "broadcast_in_dim" / "dynamic_broadcast_in_dim" – pass-through broadcast (ONNX handles implicit broadcasting).

• "call" – call to a private helper function; has extra key func.

• "reduce_max" / "reduce_sum" – reduction with keepdims; has axes.

• "skip" – shape-only op (reshape of integer tensors, concatenate of integer tensors, get_dimension_size); should be ignored by the converter.

• Any other name – direct StableHLO→ONNX op name (e.g. "sine", "add", "compare_GT").