def selective_state_update(
        cls, state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False
    ):
        """
        Argument:
            state: (batch, dim, dstate) or (batch, nheads, dim, dstate)
            x: (batch, dim) or (batch, nheads, dim)
            dt: (batch, dim) or (batch, nheads, dim)
            A: (dim, dstate) or (nheads, dim, dstate) or (dstate, dim) or (nheads, dstate, dim)
            B: (batch, dstate) or (batch, ngroups, dstate)
            C: (batch, dstate) or (batch, ngroups, dstate)
            D: (dim,) or (nheads, dim) or None
            z: (batch, dim) or (batch, nheads, dim) or None
            dt_bias: (dim,) or (nheads, dim) or None
            dt_softplus: bool
        Return:
            out: (batch, dim) or (batch, nheads, dim)
        """
        return torch.ops.torch_ipex.selective_state_update(
            state, x, dt, A, B, C, D, z, dt_bias, dt_softplus
        )

        Argument:
            state: (batch, dim, dstate) or (batch, nheads, dim, dstate)
            x: (batch, dim) or (batch, nheads, dim)
            dt: (batch, dim) or (batch, nheads, dim)
            A: (dim, dstate) or (nheads, dim, dstate) or (dstate, dim) or (nheads, dstate, dim)
            B: (batch, dstate) or (batch, ngroups, dstate)
            C: (batch, dstate) or (batch, ngroups, dstate)
            D: (dim,) or (nheads, dim) or None
            z: (batch, dim) or (batch, nheads, dim) or None
            dt_bias: (dim,) or (nheads, dim) or None
            dt_softplus: bool
        Return:
            out: (batch, dim) or (batch, nheads, dim)
        
intel_extension_for_pytorch/transformers/models/cpu/fusions/mha_fusion.py