[GOID] Use unified printing for parameter types

frontend/catalyst/decomposition/graph_op_id.py

@@ -14,7 +14,6 @@
 
 """Python implementation of Graph Operator ID."""
 
-import contextlib
 from collections.abc import Mapping, Sequence
 from typing import Any
 
@@ -29,25 +28,14 @@
     replace_wires_with_placeholder_wires,
 )
 from catalyst.from_plxpr.uid import generate_uid
-from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval
+from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval, mlir_build_context
 
 _SPECIAL_LOWERINGS = {}
 
 
-@contextlib.contextmanager
-def _attribute_context():
-    """Provide an MLIR context and location for attribute construction."""
-    if current := ir.Context.current:
-        with ir.Location.unknown(context=current):
-            yield current
-    else:
-        with ir.Context() as context, ir.Location.unknown(context=context):
-            yield context
-
-
 def format_static_data_dict_for_id(static_data):
     """Format the static-data group of a GraphOpID with MLIR's attribute printer."""
-    with _attribute_context():
+    with mlir_build_context():
         return str(get_mlir_attribute_from_pyval(static_data))
 
 

frontend/catalyst/decomposition/type_utils.py

@@ -21,10 +21,13 @@
 import jax.numpy as jnp
 import numpy as np
 import pennylane as qp
+from jax._src.interpreters.mlir import dtype_to_ir_type
 from jax._src.lib.mlir import ir
 from jax.core import ShapedArray
 from pennylane.pytrees import flatten, unflatten
 
+from catalyst.jax_extras.lowering import mlir_build_context
+
 _MLIR_DTYPES_TO_PY_DTYPES = {
     "i1": jnp.bool_,
     "i8": jnp.int8,
@@ -38,45 +41,23 @@
     "complex<f64>": jnp.complex128,
 }
 
-_PY_DTYPES_TO_MLIR_DTYPES = {v: k for k, v in _MLIR_DTYPES_TO_PY_DTYPES.items()} | {
-    float: "f64",
-    int: "i64",
-    complex: "complex<f64>",
-    ir.F16Type: "f16",
-    ir.F32Type: "f32",
-    ir.F64Type: "f64",
-    (ir.ComplexType, ir.F32Type): "complex<f32>",
-    (ir.ComplexType, ir.F64Type): "complex<f64>",
-    np.dtype("bool_"): "i1",
-    np.dtype("int8"): "i8",
-    np.dtype("int16"): "i16",
-    np.dtype("int32"): "i32",
-    np.dtype("int64"): "i64",
-    np.dtype("float16"): "f16",
-    np.dtype("float32"): "f32",
-    np.dtype("float64"): "f64",
-    np.dtype("complex64"): "complex<f32>",
-    np.dtype("complex128"): "complex<f64>",
-}
-
 
 def convert_item_to_mlir_type(item, is_special_lowering=False):
-    """Convert a string or PennyLane AbstractArray to an mlir type annotation."""
+    """Convert a string or PennyLane AbstractArray to an mlir type annotation.
+
+    The type is spelled by MLIR's own printer, applied to the type the value lowers to, which is
+    what ``printDynamicShape`` in mlir/lib/Quantum/IR/QuantumInterfaces.cpp does as well. One
+    printer spelling both sides is what keeps a rule compiled here findable by the
+    ``graph-decomposition`` pass.
+    """
     if isinstance(item, str):
         return item
 
-    if item.shape == ():
-        if is_special_lowering:
-            return _PY_DTYPES_TO_MLIR_DTYPES[item.dtype]
-        return "tensor<" + _PY_DTYPES_TO_MLIR_DTYPES[item.dtype] + ">"
-
-    return (
-        "tensor<"
-        + "x".join(str(dim_size) for dim_size in item.shape)
-        + "x"
-        + _PY_DTYPES_TO_MLIR_DTYPES[item.dtype]
-        + ">"
-    )
+    with mlir_build_context():
+        element_type = dtype_to_ir_type(np.dtype(item.dtype))
+        if is_special_lowering and item.shape == ():
+            return str(element_type)
+        return str(ir.RankedTensorType.get(item.shape, element_type))
 
 
 def get_dummy_values_for_arg(arg):

frontend/catalyst/from_plxpr/qref_operator2_primitives.py

@@ -17,6 +17,7 @@
 """
 
 # pylint: disable=unused-argument
+import numpy as np
 import pennylane as qp
 from jax._src.lib.mlir import ir
 from jax.core import ShapedArray
@@ -39,7 +40,10 @@
     inject_new_rules_into_module,
 )
 from catalyst.decomposition.graph_op_id import _SPECIAL_LOWERINGS, build_graph_op_id
-from catalyst.decomposition.type_utils import get_dummy_values_for_arg
+from catalyst.decomposition.type_utils import (
+    convert_item_to_mlir_type,
+    get_dummy_values_for_arg,
+)
 from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval
 from catalyst.jax_extras.patches import mock_attributes
 from catalyst.jax_primitives import (
@@ -318,11 +322,10 @@ def compile_decomp_rules(
     elif op_cls is qp.QubitUnitary:
         num_wires = wire_lens[0]
         matrix_size = 2**num_wires
-        dynamic_shape = {
-            qp.QubitUnitary.dynamic_argnames[0]: [
-                f"tensor<{matrix_size}x{matrix_size}xcomplex<f64>>"
-            ]
-        }
+        matrix_type = convert_item_to_mlir_type(
+            ShapedArray((matrix_size, matrix_size), np.complex128)
+        )
+        dynamic_shape = {qp.QubitUnitary.dynamic_argnames[0]: [matrix_type]}
         wire_argname = qp.QubitUnitary.wire_argnames[0]
         op_id = build_graph_op_id("QubitUnitary", dynamic_shape, {wire_argname: wire_lens[0]}, {})
 

frontend/catalyst/jax_extras/lowering.py

@@ -15,6 +15,7 @@
 
 from __future__ import annotations
 
+import contextlib
 import dataclasses
 import logging
 import textwrap
@@ -192,6 +193,17 @@ def custom_lower_jaxpr_to_module(
     return ctx.module, ctx.context
 
 
+@contextlib.contextmanager
+def mlir_build_context():
+    """Provide an MLIR context and location for attribute and type construction."""
+    if current := ir.Context.current:
+        with ir.Location.unknown(context=current):
+            yield current
+    else:
+        with ir.Context() as context, ir.Location.unknown(context=context):
+            yield context
+
+
 def get_mlir_attribute_from_pyval(value):
     """
     Given a value of any type, construct an mlir attribute of corresponding type.