    def _bmm_parameters(self, bmm_op_name, B, M, N, K):
        """
        Return a dict of parameters used for constructing bmm ops
        """
        bmm_op_name = bmm_op_name[:7]
        bmm_rcr_dict = {
            "a_shape": [B, M, K],
            "b_shape": [B, N, K],
            "c_shape": [B, M, N],
            "a_permute": None,
            "b_permute": [0, 2, 1],
        }
        bmm_crr_dict = {
            "a_shape": [B, K, M],
            "b_shape": [B, K, N],
            "c_shape": [B, M, N],
            "a_permute": [0, 2, 1],
            "b_permute": None,
        }
        bmm_ccr_dict = {
            "a_shape": [B, K, M],
            "b_shape": [B, N, K],
            "c_shape": [B, M, N],
            "a_permute": [0, 2, 1],
            "b_permute": [0, 2, 1],
        }
        bmm_rrr_dict = {
            "a_shape": [B, M, K],
            "b_shape": [B, K, N],
            "c_shape": [B, M, N],
            "a_permute": None,
            "b_permute": None,
        }
        bmm_permutes = {
            "bmm_rcr": bmm_rcr_dict,
            "bmm_crr": bmm_crr_dict,
            "bmm_ccr": bmm_ccr_dict,
            "bmm_rrr": bmm_rrr_dict,
        }
        return bmm_permutes.get(bmm_op_name)
Return a dict of parameters used for constructing bmm ops
tests/unittest/compiler/test_strided_op_cat_pattern.py