    def _infer_shapes(self):
        """infer the output shape for classic_b2b_bmm."""
        q, k, v, bias = self._attrs["inputs"]
        q_shape = q._attrs["shape"]
        k_shape = k._attrs["shape"]
        v_shape = v._attrs["shape"]
        head_dim = 2
        seq_dim = 1
        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 and len(k_shape) != 4:
            raise RuntimeError(
                f"QKV must have rank 3 or 4! 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 batch size! QKV shapes: {q_shape=}, {k_shape=}, {v_shape=}."
            )

        batch_size = q_shape[0]
        M0 = q_shape[seq_dim]
        K0 = q_shape[-1]
        if K0 != k_shape[-1]:
            raise RuntimeError(
                f"Q K shapes are not compatible! QKV shapes: {q_shape=}, {k_shape=}, {v_shape=}."
            )
        N0 = k_shape[seq_dim]
        if N0 != v_shape[seq_dim]:
            raise RuntimeError(
                f"K V shapes are not compatible! QKV shapes: {q_shape=}, {k_shape=}, {v_shape=}."
            )
        N1 = v_shape[-1]
        if N0.upper_bound() > 512 or N1.upper_bound() > 512:
            raise RuntimeError(
                f"classic_b2b_bmm only supports <=512 N0 / N1. Current length: {N0=}, {N1=}"
            )
        if not isinstance(N0, IntImm) or not isinstance(N1, IntImm):
            raise RuntimeError(
                f"classic_b2b_bmm only supports static N0 / N1. Current {N0=}, {N1=}."
            )
        if self._attrs["causal_type"] != CausalType.NO_CAUSAL:
            if M0 != N0:
                raise RuntimeError(
                    f"When causal_type is enabled, M0 must be equal to N0. Current {M0=}, {N0=}."
                )
        bias_shape = bias._attrs["shape"]

        is_multihead = len(q_shape) == 4
        if is_multihead:
            num_heads = q_shape[head_dim]

            output_shape = [batch_size, M0, num_heads, N1]
            if len(bias_shape) != 4:
                raise RuntimeError(
                    f"Was expecting 4-dimensional bias based on q dimensionality. {len(bias_shape)=} {len(q_shape)=}"
                )
            for bias_dim, expected_dim in zip(
                bias_shape, [batch_size, num_heads, M0, N0]
            ):
                if bias_dim != IntImm(1) and bias_dim != expected_dim:
                    raise RuntimeError(
                        f"bias shape is not compatible with Q K! "
                        f"QKV shapes: {q_shape=}, {num_heads=}, {k_shape=}, {v_shape=}, "
                        f"bias shapes: {bias_shape=}."
                    )
            # key sequence length is identical to last shape dim of bias tensor
            # so if it is also constant 1, it is not a real broadcast and permissible
            if bias_shape[-1] == IntImm(1) and k_shape[seq_dim] != IntImm(1):
                raise RuntimeError(
                    "classic_b2b_bmm op does not support broadcasting of last dimension of bias tensor (e.g. over sequence length of key and value ). Use the expand op to emulate this broadcast behavior if you need it."
                )
        else:
            num_heads = IntImm(1)
            self._attrs["num_heads"] = num_heads
            output_shape = [batch_size, M0, N1]
            if len(bias_shape) != 3:
                raise RuntimeError(
                    f"Was expecting 3-dimensional bias based on q dimensionality. {len(bias_shape)=} {len(q_shape)=}"
                )
            for bias_dim, expected_dim in zip(bias_shape, [batch_size, M0, N0]):
                if bias_dim != IntImm(1) and bias_dim != expected_dim:
                    raise RuntimeError(
                        f"bias shape is not compatible with Q K! "
                        f"QKV shapes: {q_shape=}, {num_heads=}, {k_shape=}, {v_shape=}, "
                        f"bias shapes: {bias_shape=}."
                    )

        return output_shape
infer the output shape for classic_b2b_bmm.
python/aitemplate/compiler/ops/b2b_bmm/classic_b2b_bmm.py