Implement the code for the minGRU neural network architecture in pure C.
Below is the reference python code:


# https://arxiv.org/abs/2410.01201v1

import torch
import torch.nn as nn
import torch.nn.functional as F


# appendix B
# https://github.com/glassroom/heinsen_sequence

def heinsen_associative_scan_log(log_coeffs, log_values):
    a_star = log_coeffs.cumsum(dim = 1)
    log_h0_plus_b_star = (log_values - a_star).logcumsumexp(dim = 1)
    log_h = a_star + log_h0_plus_b_star
    return log_h.exp()

# appendix B.3

def g(x):     return torch.where(x >= 0, x + 0.5, x.sigmoid())
def log_g(x): return torch.where(x >= 0, (F.relu(x) + 0.5).log(), -F.softplus(-x))

# log-space version of minGRU - B.3.1
# they enforce the hidden states to be positive

class minGRU(nn.Module):
    def __init__(self, d_model, d_inner):
        super().__init__()

        self.to_hidden_and_gate = nn.Linear(d_model, d_inner * 2, bias=False)
        self.to_out                         = nn.Linear(d_inner,   d_model,      bias=False)

    def forward(self, x, prev_hidden = None, return_next_prev_hidden = False):
        seq_len = x.shape[1]
        hidden, gate = self.to_hidden_and_gate(x).chunk(2, dim=-1)

        if seq_len == 1: # Sequential
            hidden = g(hidden)
            gate = gate.sigmoid()
            out = torch.lerp(prev_hidden, hidden, gate) if prev_hidden is not None else (hidden * gate)

        else: # Parallel
            log_coeffs = -F.softplus(gate)

            log_z = -F.softplus(-gate)
            log_tilde_h = log_g(hidden)
            log_values = log_z + log_tilde_h

            if prev_hidden is not None:
                log_values = torch.cat((prev_hidden.log(), log_values), dim = 1)
                log_coeffs = F.pad(log_coeffs, (0, 0, 1, 0))

            out = heinsen_associative_scan_log(log_coeffs, log_values)
            out = out[:, -seq_len:]

        prev_hidden = out[:, -1:]
        out = self.to_out(out)

        return out, prev_hidden




class RMSNorm(nn.Module):
    def __init__(self, d_model: int, eps: float=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d_model))

    def _norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        output = self._norm(x.float()).type_as(x)
        return output * self.weight





class minGRULM(nn.Module):
    def __init__(self, vocab_size, d_model, d_inner, n_layers):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, d_model)

        self.layers = nn.ModuleList([])

        for _ in range(n_layers):
            self.layers.append(nn.ModuleList([
                RMSNorm(d_model),
                minGRU(d_model, d_inner)
            ]))

        self.norm = RMSNorm(d_model)
        self.to_logits = nn.Linear(d_model, vocab_size, bias = False)

    def forward(
        self,
        x,
        return_loss = False,
        return_prev_hiddens = False,
        prev_hiddens = None
    ):

        if return_loss:
            x, labels = x[:, :-1], x[:, 1:]

        x = self.token_emb(x)

        # handle previous hiddens, for recurrent decoding
        if prev_hiddens is not None:
            x = x[:, -1:]

        next_prev_hiddens = []
        prev_hiddens = iter(prev_hiddens if prev_hiddens is not None else [])

        for norm, mingru,in self.layers:
            prev_hidden = next(prev_hiddens, None)

            min_gru_out, next_prev_hidden = mingru(
                norm(x),
                prev_hidden,
                return_next_prev_hidden = True
            )

            x = min_gru_out + x
            next_prev_hiddens.append(next_prev_hidden)


        embed = self.norm(x)
        logits = self.to_logits(embed)

        if not return_loss:
            if not return_prev_hiddens:
                return logits

            return logits, next_prev_hiddens

        loss = F.cross_entropy(
            logits.transpose(1, 2),
            labels
        )

        return loss