def _create_bmm_permute_graph(
        self,
        A_shape: Tuple[IntVar, int, int],
        B_shape: Tuple[IntVar, int, int],
        bmm_type: str,
        dtype: str,
    ) -> Tuple[Tensor, Tensor, Tensor]:
        """
        Create a graph consisting of bmm with given layout + permute021.
        """
        OP = getattr(ops, bmm_type, None)
        assert OP is not None

        A = Tensor(shape=A_shape, dtype=dtype, name="input_0", is_input=True)
        B = Tensor(shape=B_shape, dtype=dtype, name="input_1", is_input=True)
        inputs = [A, B]

        Y = OP()(*inputs)
        Y = ops.permute021()(Y)
        Y._attrs["name"] = "target_bmm_tensor"
        return A, B, Y

        Create a graph consisting of bmm with given layout + permute021.
        
tests/unittest/compiler/test_fuse_bmm_permute.py