onnxscript 0.1.0__py3-none-any.whl

onnxscript/rewriter/llama_rule_sets.py

@@ -0,0 +1,287 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+from typing import ClassVar
+
+import onnx.numpy_helper
+
+import onnxscript.ir as ir
+import onnxscript.rewriter._ir_utils as ir_utils
+import onnxscript.rewriter.no_op as no_op
+import onnxscript.rewriter.pattern as orp
+
+
+class CastIdentity(orp.RewriteRuleAsClass):
+    """Replaces ``Cast(., to=to)`` by ``Identity`` if possible."""
+
+    @classmethod
+    def pattern(cls, op, x, to):
+        return op.Cast(x, to=to)
+
+    @classmethod
+    def rewrite(cls, op, x: ir.Value, to: ir.Attr):
+        return op.Identity(x)
+
+    @classmethod
+    def check(cls, context, x, to) -> bool:
+        return x.dtype == to.value
+
+
+class CastCast(orp.RewriteRuleAsClass):
+    """Replaces ``Cast(Cast(X, ...), to=to)`` by ``Cast(X, to=to)``."""
+
+    _allowed_tensor_types: ClassVar = {
+        onnx.TensorProto.FLOAT,
+        onnx.TensorProto.FLOAT16,
+        onnx.TensorProto.BFLOAT16,
+        onnx.TensorProto.DOUBLE,
+    }
+
+    @classmethod
+    def pattern(cls, op, x, to, to_ignored):
+        return op.Cast(op.Cast(x, to=to_ignored), to=to)
+
+    @classmethod
+    def check(cls, context, x: ir.Value, to: ir.Attr, to_ignored: ir.Attr) -> bool:
+        return (
+            to.value in cls._allowed_tensor_types
+            and to_ignored.value in cls._allowed_tensor_types
+        )
+
+    @classmethod
+    def rewrite(cls, op, x: ir.Value, to: ir.Attr, to_ignored: ir.Attr):
+        return op.Cast(x, to=to)
+
+
+class ExpandIdentity(orp.RewriteRuleAsClass):
+    """Replaces ``Expand(..., shape)`` by ``Identity`` if possible."""
+
+    @classmethod
+    def pattern(cls, op, x, shape):
+        return op.Expand(x, shape)
+
+    @classmethod
+    def rewrite(cls, op, x: ir.Value, shape: ir.Value):
+        return op.Identity(x)
+
+    @classmethod
+    def check(cls, context, x, shape) -> bool:
+        if shape.const_value is None:
+            # Shape is not a constant and cannot be guessed.
+            return False
+        if (x_shape := x.shape) is None:
+            # We don't know the shape of the input
+            return False
+        return x_shape.dims == tuple(shape.const_value.numpy().tolist())
+
+
+class ReshapeReshape(orp.RewriteRuleAsClass):
+    """Replaces ``Reshape(Reshape(X, ...), shape)`` by ``Reshape(X, shape)``.
+    The pattern matches only if second reshape reshapes into a shape
+    with positive values.
+    """
+
+    @classmethod
+    def pattern(cls, op, x, shape_ignored, shape):
+        return op.Reshape(op.Reshape(x, shape_ignored), shape)
+
+    @classmethod
+    def rewrite(cls, op, x: ir.Value, shape_ignored: ir.Value, shape: ir.Value):
+        return op.Reshape(x, shape)
+
+    @classmethod
+    def check(cls, context, x, shape_ignored, shape) -> bool:
+        if shape_ignored.const_value is None or shape.const_value is None:
+            return False
+        if shape.const_value.numpy().min() <= 0:
+            return False
+        return True
+
+
+class SlicesSplit(orp.RewriteRuleAsClass):
+    """Replaces ``Slice(x, ...), Slice(x, ...)``
+    by ``Split(x, ...)`` if possible.
+    """
+
+    @classmethod
+    def pattern(cls, op, x, begin0, end0, axes0, begin1, end1, axes1):
+        return op.Slice(x, begin0, end0, axes0), op.Slice(x, begin1, end1, axes1)
+
+    @classmethod
+    def check(cls, context, x, begin0, end0, axes0, begin1, end1, axes1) -> bool:
+        if (
+            axes0.const_value is None
+            or axes1.const_value is None
+            or axes0.const_value.numpy().tolist() != axes1.const_value.numpy().tolist()
+        ):
+            return False
+        axes = axes0.const_value.numpy().tolist()
+        if len(axes) != 1:
+            return False
+        if x.shape:
+            rk = len(x.shape)
+        else:
+            rk = x.rank
+        if axes[0] != -1 and axes[0] != rk - 1:
+            return False
+        if (
+            begin0.const_value is None
+            or end0.const_value is None
+            or begin1.const_value is None
+            or end1.const_value is None
+        ):
+            return False
+        if begin0.const_value.numpy().tolist() != [0]:
+            return False
+        e0, b1, e1 = (
+            end0.const_value.numpy().tolist(),
+            begin1.const_value.numpy().tolist(),
+            end1.const_value.numpy().tolist(),
+        )
+        if e0[0] != b1[0]:
+            return False
+        shape = x.shape
+        if shape is None:
+            return False
+        last_dim = shape[-1]
+        if not isinstance(last_dim, int):
+            return False
+        if last_dim != e1[0]:
+            return False
+        if last_dim // 2 != b1[0]:
+            return False
+        return True
+
+    @classmethod
+    def rewrite(cls, op, x, begin0, end0, axes0, begin1, end1, axes1):
+        return op.Split(x, num_outputs=2, axis=-1, _outputs=2)
+
+
+class TransposeIdentity(orp.RewriteRuleAsClass):
+    """Replaces ``Transpose(. perm=perm)``
+    when the permutation is identity.
+    """
+
+    @classmethod
+    def pattern(cls, op, x, perm):
+        return op.Transpose(x, perm=perm)
+
+    @classmethod
+    def check(cls, context, x: ir.Value, perm: ir.Attr) -> bool:
+        if isinstance(perm, ir.RefAttr):
+            return False
+        if perm.type == ir.AttributeType.INTS:
+            if perm.value == list(range(len(perm.value))):
+                return True
+        return False
+
+    @classmethod
+    def rewrite(cls, op, x: ir.Value, perm: ir.Attr):
+        return op.Identity(x)
+
+
+class TransposeTranspose(orp.RewriteRuleAsClass):
+    """Replaces ``Transpose(Transpose(., perm=perm1), perm=perm2)``
+    when both permutations are inverse.
+    """
+
+    @classmethod
+    def pattern(cls, op, x, perm1, perm2):
+        return op.Transpose(op.Transpose(x, perm=perm1), perm=perm2)
+
+    @classmethod
+    def check(cls, context, x: ir.Value, perm1: ir.Attr, perm2: ir.Attr) -> bool:
+        if isinstance(perm1, ir.RefAttr) or isinstance(perm2, ir.RefAttr):
+            return False
+        return True
+
+    @classmethod
+    def _apply_transpose(cls, perm: tuple[int, ...], on: list[int]) -> list[int]:
+        assert len(perm) == len(on), "length mismatch"
+        res = [-1 for i in on]
+        for i, p in enumerate(perm):
+            res[i] = on[p]
+        return res
+
+    @classmethod
+    def _apply_transposes(
+        cls, perms: list[tuple[int, ...]], on: list[int] | None = None
+    ) -> list[int]:
+        if on is None:
+            on = list(range(len(perms[0])))
+        for p in perms:
+            on = cls._apply_transpose(p, on)
+        return on
+
+    @classmethod
+    def rewrite(cls, op, x: ir.Value, perm1: ir.Attr, perm2: ir.Attr):
+        first = list(range(len(perm1.value)))
+        last = cls._apply_transposes([perm1.value, perm2.value])
+        if first == last:
+            return op.Identity(x)
+        return op.Transpose(x, perm=last)
+
+
+class UnsqueezeUnsqueeze(orp.RewriteRuleAsClass):
+    """Replaces ``Unsqueeze(Unsqueeze(., axes1), axes2)`` with one Unsqueeze."""
+
+    @classmethod
+    def pattern(cls, op, x, axes1, axes2):
+        return op.Unsqueeze(op.Unsqueeze(x, axes1), axes2)
+
+    @classmethod
+    def rewrite(cls, op, x: ir.Value, axes1: ir.Value, axes2: ir.Value):
+        v1 = ir_utils.get_singleton_value(axes1)
+        v2 = ir_utils.get_singleton_value(axes2)
+        axes = [v1, v2] if v1 < v2 else [v2, v1 + 1]
+        return op.Unsqueeze(x, op.Constant(value=ir.tensor(axes, dtype=ir.DataType.INT64)))
+
+    @classmethod
+    def check(cls, context, x, axes1, axes2) -> bool:
+        del context  # Unused
+        del x  # Unused
+        # Currently restricted to single element positive axis
+        v1 = ir_utils.get_singleton_value(axes1)
+        v2 = ir_utils.get_singleton_value(axes2)
+        if v1 is None or v2 is None:
+            return False
+        if (v1 < 0) or (v2 < 0):
+            return False
+        return True
+
+
+cast_cast_rule = orp.make_rewrite_rule_from_class(CastCast)
+cast_identity_rule = orp.make_rewrite_rule_from_class(CastIdentity)
+expand_identity_rule = orp.make_rewrite_rule_from_class(ExpandIdentity)
+reshape_reshape_rule = orp.make_rewrite_rule_from_class(ReshapeReshape)
+slice_split_rule = orp.make_rewrite_rule_from_class(SlicesSplit, True)
+transpose_identity_rule = orp.make_rewrite_rule_from_class(TransposeIdentity)
+transpose_transpose_rule = orp.make_rewrite_rule_from_class(TransposeTranspose)
+unsqueeze_unsqueeze_rule = orp.make_rewrite_rule_from_class(UnsqueezeUnsqueeze)
+
+
+def llama_p0_rule_set() -> orp.RewriteRuleSet:
+    """Returns a set of rules which should be applied
+    before any other one as they usually remove unnecessary computation
+    such as the multiplication by 1 or two consecutive transpose.
+
+    Returns:
+        RewriteRuleSet
+    """
+    return orp.RewriteRuleSet(
+        [
+            no_op.mul_by_1_rule,
+            no_op.add_0_rule,
+            no_op.add_0_rule,
+            no_op.div_by_1_rule,
+            cast_cast_rule,
+            cast_identity_rule,
+            expand_identity_rule,
+            reshape_reshape_rule,
+            slice_split_rule,
+            transpose_identity_rule,
+            transpose_transpose_rule,
+            unsqueeze_unsqueeze_rule,
+        ]
+    )

onnxscript/rewriter/no_op.py

@@ -0,0 +1,56 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from onnxscript.rewriter import pattern
+
+# TODO: Support 1-D constant tensors
+# https://github.com/microsoft/onnx-rewriter/issues/186
+
+
+# Pattern to match against
+def mul_by_1(op, x):
+    return x * 1
+
+
+def add_0(op, x):
+    return x + 0
+
+
+def sub_0(op, x):
+    return x - 0
+
+
+def div_by_1(op, x):
+    return x / 1
+
+
+def dropout_zero(op, x):
+    return op.Dropout(x, ratio=0.0)
+
+
+def dropout_inference(op, x):
+    return op.Dropout(x, training_mode=False)
+
+
+# Replacement
+def identity(op, x, **_):
+    return op.Identity(x)
+
+
+mul_by_1_rule = pattern.RewriteRule(mul_by_1, identity)
+add_0_rule = pattern.RewriteRule(add_0, identity)
+sub_0_rule = pattern.RewriteRule(sub_0, identity)
+div_by_1_rule = pattern.RewriteRule(div_by_1, identity)
+dropout_zero_rule = pattern.RewriteRule(dropout_zero, identity)
+dropout_inference_rule = pattern.RewriteRule(dropout_inference, identity)
+# TODO: Include Mul by 0, 0 by Mul, 0 by Div? Those would be 0s, but not no-ops
+
+rules = pattern.RewriteRuleSet(
+    [
+        *mul_by_1_rule.commute(),
+        *add_0_rule.commute(),
+        sub_0_rule,
+        div_by_1_rule,
+        dropout_zero_rule,
+        dropout_inference_rule,
+    ]
+)

onnxscript/rewriter/onnxruntime/__init__.py

@@ -0,0 +1,50 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+import onnx
+
+from onnxscript.rewriter import function_rule, pattern
+from onnxscript.rewriter import rewrite as _rewrite
+from onnxscript.rewriter.onnxruntime import (
+    fused_matmul_rule_sets,
+    group_normalization_merge_silu,
+    instance_to_group_normalization,
+    softmax,
+    transformers,
+)
+
+ORT_FUNCTION_REWRITE_RULES = [*transformers.TRANSFORMERS_FUNCTION_REWRITE_RULES]
+
+ORT_PATTERN_REWRITE_RULES = [
+    *softmax.rules.rules,
+    *instance_to_group_normalization.rules.rules,
+    # NOTE: group normalization merge silu should be applied after instance to group normalization
+    *group_normalization_merge_silu.rules.rules,
+    *fused_matmul_rule_sets.fused_matmul_rule_sets(),
+]
+
+
+def rewrite(
+    model_proto: onnx.ModelProto,
+    /,
+    function_rules: list[type[function_rule.FunctionRewriteRule]] | None = None,
+    pattern_rules: list[pattern.RewriteRule] | None = None,
+) -> onnx.ModelProto:
+    """Rewrite the model using the given rules.
+
+    Args:
+        model_proto: The model to rewrite.
+        function_rules: The function rewrite rules to apply. If None, the default rules
+            for onnxruntime are used.
+        pattern_rules: The pattern rewrite rules to apply. If None, the default rules
+            for onnxruntime are used.
+
+    Returns:
+        The rewritten model.
+    """
+    function_rules = function_rules or ORT_FUNCTION_REWRITE_RULES
+    pattern_rules = pattern_rules or ORT_PATTERN_REWRITE_RULES
+    return _rewrite(
+        model_proto, function_rewrite_rules=function_rules, pattern_rewrite_rules=pattern_rules
+    )

onnxscript/rewriter/onnxruntime/bfloat16_utils/bfloat16_converter.py

@@ -0,0 +1,99 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+import logging
+
+from onnxscript import ir
+
+logger = logging.getLogger(__name__)
+
+
+def _convert_inputs_from_bfloat16_to_float16(value: ir.Value) -> None:
+    if value.dtype != ir.DataType.BFLOAT16:
+        return
+    value.dtype = ir.DataType.FLOAT16
+    _insert_cast_nodes_for_float16_to_bfloat16_to_inputs(value)
+
+
+def _convert_outputs_from_bfloat16_to_float16(value: ir.Value) -> None:
+    if value.dtype != ir.DataType.BFLOAT16:
+        return
+    _insert_cast_nodes_for_bfloat16_to_float16_to_outputs(value)
+
+
+def _insert_cast_nodes_for_float16_to_bfloat16_to_inputs(value: ir.Value) -> None:
+    user_nodes_and_indices = tuple(value.uses())
+
+    attr = ir.AttrInt64(name="to", value=ir.DataType.BFLOAT16)
+    cast = ir.Node(
+        domain="",
+        op_type="Cast",
+        inputs=[value],
+        num_outputs=1,
+        attributes=[attr],
+    )
+    cast.outputs[0].dtype = ir.DataType.BFLOAT16
+    cast.outputs[0].shape = value.shape
+
+    for node, index in tuple(value.uses()):
+        if node is cast:
+            continue
+        node.replace_input_with(index, cast.outputs[0])
+
+    # NOTE: A safer way to insert the cast node is to prepend it to the first node
+    # of the graph
+    assert user_nodes_and_indices[0][0].graph is not None, "The node should belong to a graph"
+    user_nodes_and_indices[0][0].graph[0].prepend(cast)
+
+
+def _insert_cast_nodes_for_bfloat16_to_float16_to_outputs(value: ir.Value) -> None:
+    node = value.producer()
+    index = value.index()
+    if node is None or index is None:
+        logger.warning("Output value %s has no producer or index", value)
+        return
+
+    attr = ir.AttrInt64(name="to", value=ir.DataType.FLOAT16)
+    cast = ir.Node(
+        domain="",
+        op_type="Cast",
+        inputs=[node.outputs[index]],
+        num_outputs=1,
+        attributes=[attr],
+    )
+    cast.outputs[0].dtype = ir.DataType.FLOAT16
+    cast.outputs[0].shape = node.outputs[index].shape
+    node.append(cast)
+
+    assert node.graph is not None, "Node graph should not be None"
+    # Update graph/function outputs
+    for idx, graph_or_function_output in enumerate(node.graph.outputs):
+        if graph_or_function_output == node.outputs[index]:
+            node.graph.outputs[idx] = cast.outputs[0]
+    # Swap the output name of the node with the output name of the cast node to
+    # preserve the output name in the graph
+    node.outputs[index].name, cast.outputs[0].name = (
+        cast.outputs[0].name,
+        node.outputs[index].name,
+    )
+
+
+def dtype_adapter_for_bfloat16_model(model: ir.Model) -> None:
+    """Adapt the model datatype if it's bfloat16.
+
+    Because onnxruntime does not support bfloat16 as input/output datatype, we need to
+    convert the bfloat16 datatype to float16. This function will convert the bfloat16
+    datatype to float16 and insert Cast nodes to convert float16 to bfloat16.
+
+    Model:
+        inputs(float16) -> Cast(bfloat16) -> nodes(bfloat16) -> Cast(float16) -> outputs(float16)
+
+    TODO: Delete this function after onnxruntime supports bfloat16.
+
+    Args:
+        model: The model to adapt.
+
+    """
+    for input in model.graph.inputs:
+        _convert_inputs_from_bfloat16_to_float16(input)
+    for output in model.graph.outputs:
+        _convert_outputs_from_bfloat16_to_float16(output)

onnxscript/rewriter/onnxruntime/fused_matmul_rule_sets.py

@@ -0,0 +1,177 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+from typing import ClassVar
+
+import onnxscript.rewriter.pattern as orp
+
+
+class FusedMatMulDiv1(orp.RewriteRuleAsClass):
+    """Replaces ``MatMul + Div`` by FusedMatMul."""
+
+    @classmethod
+    def pattern(cls, op, x, y, cst):
+        return op.Div(op.MatMul(x, y), cst)
+
+    @classmethod
+    def check(cls, context, x, y, cst) -> bool:
+        if cst.const_value is None:
+            return False
+        value = cst.const_value.numpy()
+        if value.size > 1:
+            return False
+        return True
+
+    @classmethod
+    def rewrite(cls, op, x, y, cst):
+        value = cst.const_value.numpy()
+        c = float(value[0] if value.shape == (1,) else value)
+        return op.FusedMatMul(x, y, alpha=1 / c, _domain="com.microsoft")
+
+
+class FusedMatMulDiv2(orp.RewriteRuleAsClass):
+    """Replaces ``FusedMatMul + Div`` by FusedMatMul."""
+
+    @classmethod
+    def pattern(cls, op, x, y, cst):
+        return op.Div(op.FusedMatMul(x, y, _domain="com.microsoft"), cst)
+
+    @classmethod
+    def check(cls, context, x, y, cst) -> bool:
+        if cst.const_value is None:
+            return False
+        if cst.const_value.numpy().size > 1:
+            return False
+        return True
+
+    @classmethod
+    def rewrite(cls, op, x, y, cst):
+        value = cst.const_value.numpy()
+        c = float(value[0] if value.shape == (1,) else value)
+        node = list(x.uses())[0][0]  # noqa: RUF015
+
+        kwargs = {}
+        alpha = node.attributes.get("alpha", None)
+        kwargs["alpha"] = alpha.value / c if alpha else 1.0 / c
+        for name in ["transA", "transB", "transBatchA", "transBatchB"]:
+            att = node.attributes.get(name)
+            if att:
+                kwargs[name] = att.value
+        return op.FusedMatMul(x, y, **kwargs, _domain="com.microsoft")
+
+
+class _TransposeMatMulBase(orp.RewriteRuleAsClass):
+    _pos: ClassVar = 1
+
+    @classmethod
+    def check(cls, context, x, y) -> bool:
+        perm = list((x if cls._pos == 1 else y).uses())[0][0].attributes["perm"].value  # noqa: RUF015
+        expected_perm = list(range(len(perm)))
+        expected_perm[-2], expected_perm[-1] = expected_perm[-1], expected_perm[-2]
+        return perm == expected_perm
+
+    @classmethod
+    def rewrite(cls, op, x, y):
+        node = list((x if cls._pos == 2 else y).uses())[0][0]  # noqa: RUF015
+        kwargs = {}
+        for name in ["alpha", "transA", "transB", "transBatchA", "transBatchB"]:
+            att = node.attributes.get(name)
+            if att:
+                kwargs[name] = att.value
+        name = "transA" if cls._pos == 1 else "transB"
+        kwargs[name] = 1 - kwargs.get(name, 0)
+        return op.FusedMatMul(x, y, **kwargs, _domain="com.microsoft")
+
+
+class TransposeMatMul1(_TransposeMatMulBase):
+    """Replaces ``Transpose + (Fused)MatMul`` by FusedMatMul."""
+
+    @classmethod
+    def pattern(cls, op, x, y):
+        return op.MatMul(op.Transpose(x), y)
+
+
+class TransposeFusedMatMul1(TransposeMatMul1):
+    """Replaces ``Transpose + (Fused)MatMul`` by FusedMatMul."""
+
+    @classmethod
+    def pattern(cls, op, x, y):
+        return op.FusedMatMul(op.Transpose(x), y, _domain="com.microsoft")
+
+
+class TransposeMatMul2(_TransposeMatMulBase):
+    """Replaces ``Transpose + (Fused)MatMul`` by FusedMatMul."""
+
+    _pos: ClassVar = 2
+
+    @classmethod
+    def pattern(cls, op, x, y):
+        return op.MatMul(x, op.Transpose(y))
+
+
+class TransposeFusedMatMul2(TransposeMatMul2):
+    """Replaces ``Transpose + (Fused)MatMul`` by FusedMatMul."""
+
+    @classmethod
+    def pattern(cls, op, x, y):
+        return op.FusedMatMul(x, op.Transpose(y), _domain="com.microsoft")
+
+
+class MatMulTranspose(orp.RewriteRuleAsClass):
+    """Replaces ``MatMul + Transpose`` by FusedMatMul."""
+
+    @classmethod
+    def pattern(cls, op, x, y):
+        return op.Transpose(op.MatMul(x, y))
+
+    @classmethod
+    def check(cls, context, x, y) -> bool:
+        matmul = list(x.uses())[0][0]  # noqa: RUF015
+        transpose = list(matmul.outputs[0].uses())[0][0]  # noqa: RUF015
+        perm = transpose.attributes["perm"].value
+        expected_perm = list(range(len(perm)))
+        expected_perm[-2], expected_perm[-1] = expected_perm[-1], expected_perm[-2]
+        return perm == expected_perm
+
+    @classmethod
+    def rewrite(cls, op, x, y):
+        node = list(x.uses())[0][0]  # noqa: RUF015
+        kwargs = {}
+        for name in ["alpha", "transA", "transB", "transBatchA", "transBatchB"]:
+            att = node.attributes.get(name)
+            if att:
+                kwargs[name] = att.value
+        for name in ["transA", "transB"]:
+            kwargs[name] = 1 - kwargs.get(name, 0)
+        return op.FusedMatMul(y, x, **kwargs, _domain="com.microsoft")
+
+
+class FusedMatMulTranspose(MatMulTranspose):
+    """Replaces ``MatMul + Transpose`` by FusedMatMul."""
+
+    @classmethod
+    def pattern(cls, op, x, y):
+        return op.Transpose(op.FusedMatMul(x, y, _domain="com.microsoft"))
+
+
+def fused_matmul_rule_sets() -> orp.RewriteRuleSet:
+    """Returns a set of rules introducting onnxruntime contrib obs.
+    This requires onnxruntime to run the model after
+    it is rewritten.
+
+    Returns:
+        RewriteRuleSet
+    """
+    return orp.RewriteRuleSet(
+        [
+            orp.make_rewrite_rule_from_class(FusedMatMulDiv1, True),
+            orp.make_rewrite_rule_from_class(FusedMatMulDiv2, True),
+            orp.make_rewrite_rule_from_class(FusedMatMulTranspose, True),
+            orp.make_rewrite_rule_from_class(MatMulTranspose, True),
+            orp.make_rewrite_rule_from_class(TransposeMatMul1, True),
+            orp.make_rewrite_rule_from_class(TransposeFusedMatMul1, True),
+            orp.make_rewrite_rule_from_class(TransposeMatMul2, True),
+            orp.make_rewrite_rule_from_class(TransposeFusedMatMul2, True),
+        ]
+    )