def selective_scan_ref(
            u,
            delta,
            A,
            B,
            C,
            D=None,
            z=None,
            delta_bias=None,
            delta_softplus=False,
            return_last_state=False,
        ):
            """
            u: r(B D L)
            delta: r(B D L)
            A: c(D N) or r(D N)
            B: c(D N) or r(B N L) or r(B N 2L) or r(B G N L) or (B G N L)
            C: c(D N) or r(B N L) or r(B N 2L) or r(B G N L) or (B G N L)
            D: r(D)
            z: r(B D L)
            delta_bias: r(D), fp32

            out: r(B D L)
            last_state (optional): r(B D dstate) or c(B D dstate)
            """
            dtype_in = u.dtype
            u = u.float()
            delta = delta.float()
            if delta_bias is not None:
                delta = delta + delta_bias[..., None].float()
            if delta_softplus:
                delta = F.softplus(delta)
            batch, dim, dstate = u.shape[0], A.shape[0], A.shape[1]
            is_variable_B = B.dim() >= 3
            is_variable_C = C.dim() >= 3
            if A.is_complex():
                if is_variable_B:
                    B = torch.view_as_complex(B.float().view(*B.shape[:-1], -1, 2))
                if is_variable_C:
                    C = torch.view_as_complex(C.float().view(*C.shape[:-1], -1, 2))
            else:
                B = B.float()
                C = C.float()
            x = A.new_zeros((batch, dim, dstate))
            ys = []
            deltaA = torch.exp(torch.einsum("bdl,dn->bdln", delta, A))
            if not is_variable_B:
                deltaB_u = torch.einsum("bdl,dn,bdl->bdln", delta, B, u)
            else:
                if B.dim() == 3:
                    deltaB_u = torch.einsum("bdl,bnl,bdl->bdln", delta, B, u)
                else:
                    B = B.repeat(1, dim // B.shape[1], 1, 1)
                    deltaB_u = torch.einsum("bdl,bdnl,bdl->bdln", delta, B, u)
            if is_variable_C and C.dim() == 4:
                C = C.repeat(1, dim // C.shape[1], 1, 1)
            last_state = None
            for i in range(u.shape[2]):
                x = deltaA[:, :, i] * x + deltaB_u[:, :, i]
                if not is_variable_C:
                    y = torch.einsum("bdn,dn->bd", x, C)
                else:
                    if C.dim() == 3:
                        y = torch.einsum("bdn,bn->bd", x, C[:, :, i])
                    else:
                        y = torch.einsum("bdn,bdn->bd", x, C[:, :, :, i])
                if i == u.shape[2] - 1:
                    last_state = x
                if y.is_complex():
                    y = y.real * 2
                ys.append(y)
            y = torch.stack(ys, dim=2)  # (batch dim L)
            out = y if D is None else y + u * D.unsqueeze(-1)
            if z is not None:
                out = out * F.silu(z)
            out = out.to(dtype=dtype_in)
            return out if not return_last_state else (out, last_state)

            u: r(B D L)
            delta: r(B D L)
            A: c(D N) or r(D N)
            B: c(D N) or r(B N L) or r(B N 2L) or r(B G N L) or (B G N L)
            C: c(D N) or r(B N L) or r(B N 2L) or r(B G N L) or (B G N L)
            D: r(D)
            z: r(B D L)
            delta_bias: r(D), fp32

            out: r(B D L)
            last_state (optional): r(B D dstate) or c(B D dstate)
            
tests/cpu/test_cpu_ops.py