    def _infer_shapes(self):
        """infer the output shape for grouped_classic_b2b_bmm."""
        q, k, v, bias = self._attrs["inputs"]
        if not (q.is_jagged() and k.is_jagged() and v.is_jagged()):
            raise RuntimeError(f"{q=}, {k=}, {v=} must be jagged!")
        q_shape = q._attrs["shape"]
        k_shape = k._attrs["shape"]
        v_shape = v._attrs["shape"]
        bias_shape = bias._attrs["shape"]
        if len(q_shape) != len(k_shape) or len(q_shape) != len(v_shape):
            raise RuntimeError(
                f"QKV ranks must be the same! QKV shapes: {q_shape=}, {k_shape=}, {v_shape=}."
            )
        if len(q_shape) != 3:
            raise RuntimeError(
                f"QKV must have rank == 3! Current rank: {len(q_shape)}, QKV shapes: {q_shape=}, {k_shape=}, {v_shape=}."
            )
        if q_shape[0] != k_shape[0] or q_shape[0] != v_shape[0]:
            raise RuntimeError(
                f"QKV must have same jagged_dim (batch_size and seq_length)! QKV shapes: {q_shape=}, {k_shape=}, {v_shape=}."
            )
        if q_shape[1] != k_shape[1] or q_shape[1] != v_shape[1]:
            raise RuntimeError(
                f"QKV must have same head size! QKV shapes: {q_shape=}, {k_shape=}, {v_shape=}."
            )

        if q_shape[2] != k_shape[2]:
            raise RuntimeError(
                f"Q and K shapes are not compatible ( inner dimension for Matmul must be identical ) - Q shape: {q_shape=}, K shape: {k_shape=}."
            )

        batch_size = q_shape[0]
        K0 = q_shape[-1]
        if K0 != k_shape[-1]:
            raise RuntimeError(
                f"Q and K shapes are not compatible! QKV shapes: {q_shape=}, {k_shape=}, {v_shape=}."
            )

        num_heads = q_shape[1]
        output_shape = [q_shape[0], num_heads, v_shape[2]]

        batch_size = q_shape[0].batch_dim()
        max_seq_len = q_shape[0].jagged_dims()[0].max_value()
        if isinstance(max_seq_len, IntVar):
            if max_seq_len.lower_bound() != max_seq_len.upper_bound():
                raise RuntimeError(
                    "Maximum sequence length needs to be a fixed (IntImm) dimension. "
                )
            max_seq_len = max_seq_len.upper_bound()

        # This is a current limitation of the classic op due to grid layout and test results
        if (
            (not _is_power_of_two(max_seq_len))
            or (max_seq_len > 512)
            or (max_seq_len < 64)
        ):
            raise RuntimeError(
                f"Maximum sequence length needs to be a fixed (IntImm) dimension with a power of two between 64 and 512 for the grouped classic b2b op to work. Actual value: {max_seq_len=}. {type(max_seq_len)=}"
            )
        if len(bias_shape) != 4:
            raise RuntimeError(f"Expected bias rank 4. Current bias rank: {len(bias)}.")

        bias_expected_shape = [
            batch_size,
            num_heads,
            max_seq_len,
            max_seq_len,
        ]
        broadcastable, _ = shape_utils.get_broadcast_max_shape(
            bias_shape, bias_expected_shape
        )
        if not broadcastable:
            raise RuntimeError(
                f"bias shape is not compatible with Q K! "
                f"QKV shapes: {q_shape=}, {k_shape=}, {v_shape=}, "
                f"bias shapes: {bias_shape=}, {bias_expected_shape=}."
            )
        if bias_shape[-1] != bias_expected_shape[-1]:
            raise RuntimeError(
                f"Bias last dim is not broadcastable! Expected shape: {bias_expected_shape[-1]}, current bias shape: {bias_shape}"
            )
        return output_shape, max_seq_len
infer the output shape for grouped_classic_b2b_bmm.
python/aitemplate/compiler/ops/b2b_bmm/grouped_classic_b2b_bmm.py