def __call__(self, a: Tensor, b: Tensor) -> Tensor:
        """Call bmm_rrr_permute with tensors a, b

        Parameters
        ----------
        a : Tensor
            Tensor in shape (B, M, K)
        b : Tensor
            Tensor in shape (B, K, N)

        Returns
        -------
        Tensor
            Tensor in shape (B // D1, M, D1, N) for 0213 permute

        Raises
        ------
        NotImplementedError
            Permute layout not implemented yet
        """
        a, b = self._align_ab(a, b)
        self._attrs["inputs"] = [a, b]
        self._attrs["input_accessors"] = [TensorAccessor(a), TensorAccessor(b)]
        self._set_depth()
        self._sanity_check(a, b)
        output_shape = self._infer_shapes(a, b)

        output = Tensor(output_shape, src_ops={self}, dtype=a.dtype())
        self._attrs["outputs"] = [output]
        self._attrs["output_accessors"] = [TensorAccessor(output)]

        if self._attrs["layout"] == "Permute4DBMM_0213":
            b, m, n = output_shape
            d1 = self._attrs["shape"][0]
            output_shape = [-1, m, d1, n]
            self._extract_epilogue_alignment(output_shape)
            return reshape()(output, output_shape)
        else:
            raise NotImplementedError(
                "{} is not implemented!".format(self._attrs["layout"])
            )
Call bmm_rrr_permute with tensors a, b

        Parameters
        ----------
        a : Tensor
            Tensor in shape (B, M, K)
        b : Tensor
            Tensor in shape (B, K, N)

        Returns
        -------
        Tensor
            Tensor in shape (B // D1, M, D1, N) for 0213 permute

        Raises
        ------
        NotImplementedError
            Permute layout not implemented yet
        
python/aitemplate/compiler/ops/gemm_universal/bmm_rrr_permute.py