def selective_scan_fn(
        cls,
        u,
        delta,
        A,
        B,
        C,
        D=None,
        z=None,
        delta_bias=None,
        delta_softplus=False,
        return_last_state=False,
    ):
        """
        u: (B D L) or (B L D)
        delta: same shape as u
        A: (D N) or (N D)
        B: (B N L) or (B N 2L) or (B G N L)
        C: (B N L) or (B N 2L) or (B G N L)
        D: (D) or None
        z: (B D L) or None
        delta_bias: (D) or None, fp32

        out: (B D L)
        last_state (optional): (B D dstate)
        """
        out, ssm_state = torch.ops.torch_ipex.selective_scan_fn(
            u,
            delta.to(A.dtype),
            A,
            B,
            C,
            D,
            z,
            delta_bias,
            delta_softplus,
            return_last_state,
        )
        return out if not return_last_state else (out, ssm_state)

        u: (B D L) or (B L D)
        delta: same shape as u
        A: (D N) or (N D)
        B: (B N L) or (B N 2L) or (B G N L)
        C: (B N L) or (B N 2L) or (B G N L)
        D: (D) or None
        z: (B D L) or None
        delta_bias: (D) or None, fp32

        out: (B D L)
        last_state (optional): (B D dstate)
        
intel_extension_for_pytorch/transformers/models/cpu/fusions/mha_fusion.py