def selective_state_update_ref(
            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)
                B: (batch, dstate) or (batch, ngroups, dstate)
                C: (batch, dstate) or (batch, ngroups, dstate)
                D: (dim,) or (nheads, dim)
                z: (batch, dim) or (batch, nheads, dim)
                dt_bias: (dim,) or (nheads, dim)
            Return:
                out: (batch, dim) or (batch, nheads, dim)
            """
            has_heads = state.dim() > 3
            if state.dim() == 3:
                state = state.unsqueeze(1)
            if x.dim() == 2:
                x = x.unsqueeze(1)
            if dt.dim() == 2:
                dt = dt.unsqueeze(1)
            if A.dim() == 2:
                A = A.unsqueeze(0)
            if B.dim() == 2:
                B = B.unsqueeze(1)
            if C.dim() == 2:
                C = C.unsqueeze(1)
            if D is not None and D.dim() == 1:
                D = D.unsqueeze(0)
            if z is not None and z.dim() == 2:
                z = z.unsqueeze(1)
            if dt_bias is not None and dt_bias.dim() == 1:
                dt_bias = dt_bias.unsqueeze(0)
            batch, nheads, dim, dstate = state.shape
            assert x.shape == (batch, nheads, dim)
            assert dt.shape == x.shape
            assert A.shape == (nheads, dim, dstate)
            ngroups = B.shape[1]
            assert nheads % ngroups == 0, "nheads must be divisible by ngroups"
            assert B.shape == (batch, ngroups, dstate)
            assert C.shape == B.shape
            if D is not None:
                assert D.shape == (nheads, dim)
            if z is not None:
                assert z.shape == x.shape
            if dt_bias is not None:
                assert dt_bias.shape == (nheads, dim)
                dt = dt + dt_bias
            dt = F.softplus(dt) if dt_softplus else dt
            dA = torch.exp(
                rearrange(dt, "b h d -> b h d 1") * A
            )  # (batch, nheads, dim, dstate)
            B = repeat(
                B, "b g n -> b (g h) n", h=nheads // ngroups
            )  # (batch, nheads, dstate)
            C = repeat(
                C, "b g n -> b (g h) n", h=nheads // ngroups
            )  # (batch, nheads, dstate)
            dB = rearrange(dt, "b h d -> b h d 1") * rearrange(
                B, "b h n -> b h 1 n"
            )  # (batch, nheads, dim, dstate)
            state.copy_(
                state * dA + dB * rearrange(x, "b h d -> b h d 1")
            )  # (batch, dim, dstate
            out = torch.einsum("bhdn,bhn->bhd", state.to(C.dtype), C)
            if D is not None:
                out += (x * D).to(out.dtype)
            out = (out if z is None else out * F.silu(z)).to(x.dtype)
            if not has_heads:
                out = out.squeeze(1)
            return out

            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)
                B: (batch, dstate) or (batch, ngroups, dstate)
                C: (batch, dstate) or (batch, ngroups, dstate)
                D: (dim,) or (nheads, dim)
                z: (batch, dim) or (batch, nheads, dim)
                dt_bias: (dim,) or (nheads, dim)
            Return:
                out: (batch, dim) or (batch, nheads, dim)
            
tests/cpu/test_ipex_llm_module.py