    def _infer_shapes(self, a: Tensor, b: Tensor) -> List[IntVar]:
        """Given input tensors, infers output tensor shapes."""

        a_shapes = a._attrs["shape"]
        if len(a_shapes) != 3:
            raise RuntimeError(
                "bmm operand A should have 3 dimensions! Current shape: {}.".format(
                    a_shapes
                )
            )
        b_shapes = b._attrs["shape"]
        if len(b_shapes) != 3:
            raise RuntimeError(
                "bmm operand B should have 3 dimensions! Current shape: {}.".format(
                    b_shapes
                )
            )
        batch_size_a = a_shapes[0]
        batch_size_b = b_shapes[0]
        if batch_size_a != batch_size_b:
            raise RuntimeError(
                "bmm operand A and B should have same batch_size! "
                "Current shape A: {} shape B: {} .".format(a_shapes, b_shapes)
            )
        assert (
            a_shapes[2] == b_shapes[1]
        ), f"bmm operand A and B should have same K dim! Current shape A: {a_shapes}, shape B: {b_shapes}"
        m_values = a_shapes[1]._attrs["values"]
        # TODO: remove shape check after fixing the kernel
        assert all(
            val % 8 == 0 for val in m_values
        ), f"M should be multiples of 8. M: {a_shapes[1]}"
        n_values = b_shapes[2]._attrs["values"]
        assert all(
            val % 8 == 0 for val in n_values
        ), f"N should be multiples of 8. N: {b_shapes[2]}"
        c_shapes = [batch_size_a, a_shapes[1], b_shapes[2]]
        return c_shapes
Given input tensors, infers output tensor shapes.
python/aitemplate/compiler/ops/gemm_special/bmm_rrr_k1_tanh.py