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_ir_module(mlir_module) → List[XlaLayer]

Parses an MLIR ir.Module using Python bindings (JAX 0.10+).

Walks MLIR operations directly without converting the module to a text string, so loc(...) annotations are not required. Returns the same list of XlaLayer objects as parse_mlir().

Must be called while an active MLIR context is open (e.g., inside a with make_ir_context(): block).

Parameters:

mlir_module – an ir.Module obtained from jax.extend.mlir.deserialize_portable_artifact.

Returns:

List of XlaLayer objects in the same format as parse_mlir().

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

Parse a StableHLO MLIR module text into a list of XlaLayer objects.

Each XlaLayer has at minimum the fields 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.

• "convert" – stablehlo.convert; type cast; has shape set to the target tensor type string (e.g. "tensor<f32>").

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