#include <stdio.h>
#include <assert.h>
#include "ATen/ATen.h"

typedef at::BFloat16 bf16;

template <typename F>
__global__ void kernel_forward(const int B, const int T, const int C, const int H,
                               const F *__restrict__ const _r, const F *__restrict__ const _w, const F *__restrict__ const _k, const F *__restrict__ const _v, const F *__restrict__ const _a, const F *__restrict__ const _b,
                               F *__restrict__ const _y, float *__restrict__ const _saa, float *__restrict__ const _sss)
{
    const int e = blockIdx.x / H;
    const int h = blockIdx.x % H;
    const int i = threadIdx.x;

    float state[_N_] = {0};
    __shared__ float r[_N_], k[_N_], w[_N_], a[_N_], b[_N_];

    float v[_T_];
    for (int _t = 0; _t < T; _t++)
    {
        const int t = e*T*C + h*_N_ + i + _t * C;
        v[_t] = float(_v[t]);
    }

    for (int _t = 0; _t < T; _t++)
    {
        const int t = e*T*C + h*_N_ + i + _t * C;
        __syncthreads();
        r[i] = float(_r[t]);
        w[i] = __expf(-__expf(float(_w[t])));
        k[i] = float(_k[t]);
        a[i] = float(_a[t]);
        b[i] = float(_b[t]);
        __syncthreads();

        float sa = 0;
        #pragma unroll
        for (int j = 0; j < _N_; j++)
        {
            sa += a[j] * state[j];
        }
        _saa[t] = float(sa);

        float vv = v[_t];
        float y = 0;
        #pragma unroll
        for (int j = 0; j < _N_; j++)
        {
            float& s = state[j];
            s = s * w[j] + k[j] * vv + sa * b[j];
            y += s * r[j];
        }
        _y[t] = F(y);

        if ((_t+1) % _CHUNK_LEN_ == 0)
        {
            const int a = _t / _CHUNK_LEN_;
            const int c = _T_ / _CHUNK_LEN_;
            const int p = e*C*_N_*c + h*_N_*_N_*c + a*_N_*_N_ + i;
            #pragma unroll
            for (int j = 0; j < _N_; j++)
            {
                _sss[p + j*_N_] = float(state[j]);
            }
        }
    }
}

template <typename F>
__global__ void kernel_backward_zzz(const int B, const int T, const int C, const int H,
    const F *__restrict__ const _r, const F *__restrict__ const _w, const F *__restrict__ const _k, const F *__restrict__ const _a, const F *__restrict__ const _b, const F *__restrict__ const _gy,
    float *__restrict__ const _zzz)
{
    const int e = blockIdx.x / H;
    const int h = blockIdx.x % H;
    const int i = threadIdx.x;

    __shared__ float r[_N_], w[_N_], a[_N_], b[_N_];

    const int T_1 = e*T*C + (T-1)*C + h*_N_;
    float z[_N_];
    const float gy = _gy[T_1 + i];
    __syncthreads();
    r[i] = float(_r[T_1+i]);
    __syncthreads();
    #pragma unroll
    for (int j = 0; j < _N_; j++) 
    {
        z[j] = gy * r[j];
    }

    for (int _t = T-2; _t > _CHUNK_LEN_-1; _t--)
    {
        const int t = e*T*C + h*_N_ + _t * C + i;
        const float gy = _gy[t];
        __syncthreads();
        r[i] = float(_r[t]);
        w[i] = __expf(-__expf(float(_w[t+C])));
        a[i] = float(_a[t+C]);
        b[i] = float(_b[t+C]);
        __syncthreads();

        float zz = 0;
        #pragma unroll
        for (int j = 0; j < _N_; j++)
        {
            zz += b[j] * z[j];
        }
        #pragma unroll
        for (int j = 0; j < _N_; j++)
        {
            z[j] = z[j] * w[j] + gy * r[j] + a[j] * zz;
            // printf("t %d i %d j %d z %f\n", _t, i, j, z[j]);
        }
        if (_t % _CHUNK_LEN_ == 0)
        {
            const int a = _t / _CHUNK_LEN_ - 1;
            const int c = _T_ / _CHUNK_LEN_ - 1;
            const int p = e*C*_N_*c + h*_N_*_N_*c + a*_N_*_N_ + i;
            #pragma unroll
            for (int j = 0; j < _N_; j++)
            {
                _zzz[p + j*_N_] = float(z[j]);
            }
        }
    }
}

template <typename F>
__global__ void kernel_backward_rwkv(const int B, const int T, const int C, const int H,
    const F *__restrict__ const _r, const F *__restrict__ const _w, const F *__restrict__ const _k, const F *__restrict__ const _v, const F *__restrict__ const _a, const F *__restrict__ const _b, const float *__restrict__ const _saa, const float *__restrict__ const _sss, const float *__restrict__ const _zzz,
    const F *__restrict__ const _gy, F *__restrict__ const _gr, F *__restrict__ const _gw, F *__restrict__ const _gk, F *__restrict__ const _gv, F *__restrict__ const _ga, F *__restrict__ const _gb)
{
    const int e = blockIdx.x / H;
    const int h = blockIdx.x % H;
    const int chunk = threadIdx.x;
    const int n_chunk = _T_ / _CHUNK_LEN_;
    
    float zzz[_N_*_N_] = {0}, sss[_N_*_N_] = {999}, saa[_N_] = {999};
    float r[_N_] = {999}, w[_N_] = {0}, w1[_N_] = {999}, winv[_N_] = {999}, ww[_N_] = {999};
    float k[_N_] = {0}, v[_N_] = {999}, a[_N_] = {0}, a1[_N_] = {999}, b[_N_] = {0}, b1[_N_] = {999}, gy[_N_] = {999};

    if (chunk != n_chunk - 1)
    {
        const int p = e*T*C + (chunk+1)*_CHUNK_LEN_*C + h*_N_;
        for (int i = 0; i < _N_; i++) 
        {
            k[i] = float(_k[p+i]);
            a[i] = float(_a[p+i]);
            b[i] = float(_b[p+i]);
            w[i] = __expf(-__expf(float(_w[p+i])));
            const int p = e*C*_N_*(n_chunk-1) + h*_N_*_N_*(n_chunk-1) + chunk*_N_*_N_ + i*_N_;
            #pragma unroll
            for (int j = 0; j < _N_; j++)
            {
                zzz[i*_N_+j] = float(_zzz[p+j]);
            }
        }
    }
    for (int i = 0; i < _N_; i++)
    {
        const int p = e*C*_N_*n_chunk + h*_N_*_N_*n_chunk + chunk*_N_*_N_ + i*_N_;
        #pragma unroll
        for (int j = 0; j < _N_; j++)
        {
            sss[i*_N_+j] = float(_sss[p+j]);
        }
    }

    for (int _t = _CHUNK_LEN_-1; _t > -1; _t--)
    {
        const int t = chunk * _CHUNK_LEN_ + _t;
        const int b_t_h = e*T*C + t*C + h*_N_;
        #pragma unroll
        for (int n = 0; n < _N_; n++)
        {
            w1[n] = w[n];
            a1[n] = a[n];
            b1[n] = b[n];

            r[n] = float(_r[b_t_h+n]);
            k[n] = float(_k[b_t_h+n]);
            v[n] = float(_v[b_t_h+n]);
            a[n] = float(_a[b_t_h+n]);
            b[n] = float(_b[b_t_h+n]);
            gy[n] = float(_gy[b_t_h+n]);
            saa[n] = float(_saa[b_t_h+n]);

            ww[n] = -__expf(float(_w[b_t_h+n]));
            w[n] = __expf(ww[n]);
            ww[n] = ww[n] * w[n];
            winv[n] = 1.0 / w[n];
        }

        for (int j = 0; j < _N_; j++)
        {
            float zz = 0;
            #pragma unroll
            for (int i = 0; i < _N_; i++)
            {
                zz += b1[i] * zzz[i*_N_+j];
            }
            const float gyj = gy[j];
            #pragma unroll
            for (int i = 0; i < _N_; i++)
            {
                zzz[i*_N_+j] = zzz[i*_N_+j] * w1[i] + gyj * r[i] + a1[i] * zz;
                // printf("t %d i %d j %d z %f\n",t,i,j,zzz[i*_N_+j]);
                // printf("t %d i %d j %d s %f\n",t,i,j,sss[i*_N_+j]);
            }
        }

        for (int i = 0; i < _N_; i++)
        {
            float gr = 0;
            #pragma unroll
            for (int j = 0; j < _N_; j++)
            {
                gr += gy[j] * sss[i*_N_+j];
            }
            _gr[b_t_h+i] = F(gr);
        }

        for (int i = 0; i < _N_; i++)
        {
            const float ki = k[i];
            const float bi = b[i];
            const float wi = winv[i];
            #pragma unroll
            for (int j = 0; j < _N_; j++)
            {
                sss[i*_N_+j] = (sss[i*_N_+j] - ki * v[j] - bi * saa[j]) * wi;
            }
        }

        float gv[_N_] = {0}; float as[_N_] = {0}; float bz[_N_] = {0};
        for (int i = 0; i < _N_; i++)
        {
            const float ki = k[i];
            const float ai = a[i];
            const float bi = b[i];
            float gw = 0;
            float gk = 0;
            #pragma unroll
            for (int j = 0; j < _N_; j++)
            {
                const float sij = sss[i*_N_+j];
                const float zij = zzz[i*_N_+j];
                gv[j] += ki * zij;
                as[j] += ai * sij;
                bz[j] += bi * zij;
                gw += sij * zij;
                gk += v[j] * zij;
            }
            _gw[b_t_h+i] = F(gw * ww[i]);
            _gk[b_t_h+i] = F(gk);
        }
        for (int i = 0; i < _N_; i++)
        {
            _gv[b_t_h+i] = F(gv[i]);
            float ga = 0;
            float gb = 0;
            #pragma unroll
            for (int j = 0; j < _N_; j++)
            {
                ga += bz[j] * sss[i*_N_+j];
                gb += as[j] * zzz[i*_N_+j];
            }
            _ga[b_t_h+i] = F(ga);
            _gb[b_t_h+i] = F(gb);
        }
    }
}

void cuda_forward(int B, int T, int C, int H, bf16 *r, bf16* w, bf16 *k, bf16 *v, bf16 *a, bf16 *b, bf16 *y, float *saa, float* sss)
{
    assert(H*_N_ == C);
    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, w, k, v, a, b, y, saa, sss);
}

void cuda_backward(int B, int T, int C, int H, bf16 *r, bf16* w, bf16 *k, bf16 *v, bf16 *a, bf16 *b, float *saa, float* sss, float* zzz, bf16 *gy, bf16 *gr, bf16 *gw, bf16 *gk, bf16 *gv, bf16 *ga, bf16 *gb)
{
    assert(H*_N_ == C);
    assert(T%_CHUNK_LEN_ == 0);

    kernel_backward_zzz<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, w, k, a, b, gy, zzz);
    kernel_backward_rwkv<<<dim3(B * H), dim3(T/_CHUNK_LEN_)>>>(B, T, C, H, r, w, k, v, a, b, saa, sss, zzz, gy, gr, gw, gk, gv, ga, gb);
}
CHUNK_LEN = 16

    load(name="wkv7g", sources=["rwkv_cuda/wkv7g_op.cpp", f"rwkv_cuda/wkv7g_v1.cu"], is_python_module=False,
                        verbose=True, extra_cuda_cflags=["-res-usage", "--use_fast_math", "-O3", "-Xptxas -O3", "--extra-device-vectorization", f"-D_N_={HEAD_SIZE}", f"-D_T_={T}", f"-D_CHUNK_LEN_={CHUNK_LEN}"])
    class WKV_7g(torch.autograd.Function):
        @staticmethod
        def forward(ctx, r, w, k, v, a, b):
            with torch.no_grad():
                B, T, C = r.size()
                H = C // HEAD_SIZE
                N = HEAD_SIZE
                A = T // CHUNK_LEN
                assert HEAD_SIZE == C // H
                assert T % CHUNK_LEN == 0
                assert r.dtype == DTYPE
                assert w.dtype == DTYPE
                assert k.dtype == DTYPE
                assert v.dtype == DTYPE
                assert a.dtype == DTYPE
                assert b.dtype == DTYPE
                assert r.is_contiguous()
                assert w.is_contiguous()
                assert k.is_contiguous()
                assert v.is_contiguous()
                assert a.is_contiguous()
                assert b.is_contiguous()
                ctx.B = B
                ctx.T = T
                ctx.C = C
                ctx.H = H
                y = torch.empty((B, T, C), device=k.device, dtype=DTYPE, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
                saa = torch.empty((B, T, H, N), device=k.device, dtype=torch.float, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
                sss = torch.empty((B, H, A, N, N), device=k.device, dtype=torch.float, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
                torch.ops.wkv7g.forward(B, T, C, H, r, w, k, v, a, b, y, saa, sss)
                ctx.save_for_backward(r, w, k, v, a, b, saa, sss)
                return y
        @staticmethod
        def backward(ctx, gy):
            with torch.no_grad():
                N = HEAD_SIZE
                B = ctx.B
                T = ctx.T
                C = ctx.C
                H = ctx.H
                A = T // CHUNK_LEN
                assert gy.dtype == DTYPE
                assert gy.is_contiguous()
                r, w, k, v, a, b, saa, sss = ctx.saved_tensors
                gr = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=DTYPE, memory_format=torch.contiguous_format)#.zero_()#.uniform_(-100, 100)
                gw = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=DTYPE, memory_format=torch.contiguous_format)#.zero_()#.uniform_(-100, 100)
                gk = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=DTYPE, memory_format=torch.contiguous_format)#.zero_()#.uniform_(-100, 100)
                gv = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=DTYPE, memory_format=torch.contiguous_format)#.zero_()#.uniform_(-100, 100)
                ga = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=DTYPE, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
                gb = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=DTYPE, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
                zzz = torch.empty((B, H, A-1, N, N), device=gy.device, dtype=XTYPE, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
                torch.ops.wkv7g.backward(B, T, C, H, r, w, k, v, a, b, saa, sss, zzz, gy, gr, gw, gk, gv, ga, gb)
                del saa
                del sss
                del zzz
                return (gr, gw, gk, gv, ga, gb)
    @torch.compiler.disable
    def RUN_CUDA_RWKV7g(r, w, k, v, a, b):
        return WKV_7g.apply(r, w, k, v, a, b)

class RWKV7(nn.Module):
    def __init__(self, args, layer_id):
        super().__init__()
        self.args = args
        self.layer_id = layer_id
        self.n_embd = args.n_embd
        args.dim_att = args.n_embd

        self.head_size = HEAD_SIZE
        self.n_head = args.dim_att // self.head_size
        assert args.dim_att % self.n_head == 0

        with torch.no_grad():
            ratio_0_to_1 = layer_id / (args.n_layer - 1)  # 0 to 1
            ratio_1_to_almost0 = 1.0 - (layer_id / args.n_layer)  # 1 to ~0
            ddd = torch.ones(1, 1, args.n_embd)
            for i in range(args.n_embd):
                ddd[0, 0, i] = i / args.n_embd

            # initialization comes from fitting my RWKV-6 7B runs
            # merging r&g w&a to save params
            self.time_maa_x = nn.Parameter(1.0 - torch.pow(ddd, 0.6 * ratio_1_to_almost0 ** 0.9))
            self.time_maa_rg = nn.Parameter(1.0 - torch.pow(ddd, 0.2 * ratio_1_to_almost0))
            self.time_maa_wa = nn.Parameter(1.0 - torch.pow(ddd, 0.9 * ratio_1_to_almost0))
            self.time_maa_k = nn.Parameter(1.0 - (torch.pow(ddd, 0.9 * ratio_1_to_almost0) + 0.4 * ratio_0_to_1))
            self.time_maa_v = nn.Parameter(1.0 - (torch.pow(ddd, 0.4 * ratio_1_to_almost0) + 0.6 * ratio_0_to_1))

            decay_speed = torch.ones(args.dim_att)
            for n in range(args.dim_att):
                decay_speed[n] = -7 + 5 * (n / (args.dim_att - 1)) ** (0.85 + 1.0 * ratio_0_to_1 ** 0.5)
            self.time_decay = nn.Parameter(decay_speed.reshape(1,1,args.dim_att) + 0.5) # !!! 0.5 comes from F.softplus !!!

            self.time_faaaa = nn.Parameter(torch.zeros(1,1,self.n_head,self.head_size))
            self.time_aaaaa = nn.Parameter(torch.zeros(1,1,args.dim_att))

            def ortho_init(x, scale):
                with torch.no_grad():
                    shape = x.shape
                    if len(shape) == 2:
                        gain = math.sqrt(shape[0] / shape[1]) if shape[0] > shape[1] else 1
                        nn.init.orthogonal_(x, gain=gain * scale)
                    elif len(shape) == 3:
                        gain = math.sqrt(shape[1] / shape[2]) if shape[1] > shape[2] else 1
                        for i in range(shape[0]):
                            nn.init.orthogonal_(x[i], gain=gain * scale)
                    else:
                        assert False
                    return x

            D_MIX_LORA = 32
            self.time_maa_w1 = nn.Parameter(torch.zeros(args.n_embd, D_MIX_LORA*4))
            self.time_maa_w2 = nn.Parameter(ortho_init(torch.zeros(4, D_MIX_LORA, args.n_embd), 0.1))

            D_DECAY_LORA = 64
            self.time_decay_w1 = nn.Parameter(torch.zeros(args.n_embd, D_DECAY_LORA))
            self.time_decay_w2 = nn.Parameter(ortho_init(torch.zeros(D_DECAY_LORA, args.dim_att), 0.1))

            D_AAA_LORA = 16
            self.time_aaa_w1 = nn.Parameter(torch.zeros(args.n_embd, D_AAA_LORA))
            self.time_aaa_w2 = nn.Parameter(ortho_init(torch.zeros(D_AAA_LORA, args.dim_att), 0.1))

            D_KKK_LORA = 16
            self.time_kkk_w1 = nn.Parameter(torch.zeros(args.n_embd, D_KKK_LORA))
            self.time_kkk_w2 = nn.Parameter(ortho_init(torch.zeros(D_KKK_LORA, args.dim_att), 0.1))

            D_GATE_LORA = 128
            self.gate_w1 = nn.Parameter(ortho_init(torch.zeros(args.n_embd, D_GATE_LORA), 0.1))
            self.gate_w2 = nn.Parameter(ortho_init(torch.zeros(D_GATE_LORA, args.dim_att), 0.1))

            D_MA_LORA = 16
            self.ma_w1 = nn.Parameter(torch.zeros(args.n_embd, D_MA_LORA))
            self.ma_w2 = nn.Parameter(ortho_init(torch.zeros(D_MA_LORA, args.dim_att), 0.1))
            self.time_misc_a = nn.Parameter(torch.zeros(1,1,args.n_embd))
            D_MK_LORA = 16
            self.mk_w1 = nn.Parameter(torch.zeros(args.n_embd, D_MK_LORA))
            self.mk_w2 = nn.Parameter(ortho_init(torch.zeros(D_MK_LORA, args.dim_att), 0.1))
            self.time_misc_k = nn.Parameter(torch.zeros(1,1,args.n_embd))

            self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
            self.receptance = nn.Linear(args.n_embd, args.dim_att, bias=False)
            self.key = nn.Linear(args.n_embd, args.dim_att, bias=False)
            self.value = nn.Linear(args.n_embd, args.dim_att, bias=False)
            self.output = nn.Linear(args.dim_att, args.n_embd, bias=False)
            self.ln_x = nn.GroupNorm(self.n_head, args.dim_att, eps=64e-5)

            self.receptance.weight.data.uniform_(-0.5/(self.n_embd**0.5), 0.5/(self.n_embd**0.5))
            self.key.weight.data.uniform_(-0.05/(self.n_embd**0.5), 0.05/(self.n_embd**0.5))
            self.value.weight.data.uniform_(-0.5/(self.n_embd**0.5), 0.5/(self.n_embd**0.5))
            self.output.weight.data.zero_()

    def forward(self, x):
        B, T, C = x.size()
        H = self.n_head
        xx = self.time_shift(x) - x

        xxx = x + xx * self.time_maa_x
        xxx = torch.tanh(xxx @ self.time_maa_w1).view(B*T, 4, -1).transpose(0, 1)
        xxx = torch.bmm(xxx, self.time_maa_w2).view(4, B, T, -1)
        mrg, mwa, mk, mv = xxx.unbind(dim=0)

        xrg = x + xx * (self.time_maa_rg + mrg)
        xwa = x + xx * (self.time_maa_wa + mwa)
        xk = x + xx * (self.time_maa_k + mk)
        xv = x + xx * (self.time_maa_v + mv)

        r = self.receptance(xrg)
        w = -F.softplus(-(self.time_decay + torch.tanh(xwa @ self.time_decay_w1) @ self.time_decay_w2)) - 0.5
        k = self.key(xk)
        v = self.value(xv)
        g = torch.tanh(xrg @ self.gate_w1) @ self.gate_w2

        kk = k + torch.tanh(xk @ self.time_kkk_w1) @ self.time_kkk_w2
        kk = F.normalize(kk.view(B,T,H,-1), dim=-1, p=2.0).view(B,T,C)
        a = torch.sigmoid( self.time_aaaaa + (xwa @ self.time_aaa_w1) @ self.time_aaa_w2 )

        ma = torch.sigmoid(self.time_misc_a + (xwa @ self.ma_w1) @ self.ma_w2)
        k = k * ma + k*a * (1 - ma)
        mk = torch.sigmoid(self.time_misc_k + (xk @ self.mk_w1) @ self.mk_w2)
        k = k * torch.clamp(w*mk, max=0).exp()

        x = RUN_CUDA_RWKV7g(r.bfloat16(), w.bfloat16(), k.bfloat16(), v.bfloat16(), -kk.bfloat16(), (kk*a).bfloat16())
        x = self.ln_x(x.view(B * T, C)).view(B, T, C)

        x = x + ((r.view(B,T,H,-1)*k.view(B,T,H,-1)*self.time_faaaa).sum(dim=-1, keepdim=True) * v.view(B,T,H,-1)).view(B,T,C)
        x = self.output(x * g)
        return x

class MLP(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.c_fc    = nn.Linear(config.n_embd, 7 * config.n_embd // 2, bias=False)
        self.c_proj  = nn.Linear(7 * config.n_embd // 2, config.n_embd, bias=False)
        self.c_proj.weight.data.zero_() # zero init suggested by @Grad62304977

    def forward(self, x):
        x = self.c_fc(x)
        x = F.relu(x).square() # https://arxiv.org/abs/2109.08668v2; ~1-2% better than GELU; suggested by @SKYLINEZ007 and @Grad62304977
        x = self.c_proj(x)
        return x

class Block(nn.Module):

    def __init__(self, config, layer_id):
        super().__init__()
        self.attn = RWKV7(config, layer_id)
        self.mlp = MLP(config)

        self.ln1 = nn.LayerNorm(config.n_embd)
        self.ln2 = nn.LayerNorm(config.n_embd)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x



Update the below tensorflow model to match the above pytorch model

import tensorflow as tf
from tensorflow.keras import layers
from dataclasses import dataclass
import math

HEAD_SIZE = 64  # Define according to your model

class RWKV7(tf.keras.layers.Layer):
    def __init__(self, config, layer_id, **kwargs):
        super(RWKV7, self).__init__(**kwargs)
        self.config = config
        self.layer_id = layer_id
        self.n_embd = config.n_embd
        self.dim_att = config.n_embd

        self.head_size = HEAD_SIZE
        self.n_head = self.dim_att // self.head_size
        assert self.dim_att % self.n_head == 0, "dim_att must be divisible by n_head"

        # Initialize ratio parameters
        ratio_0_to_1 = layer_id / (config.n_layer - 1)
        ratio_1_to_almost0 = 1.0 - (layer_id / config.n_layer)

        # Initialize ddd similar to PyTorch
        ddd = tf.linspace(0.0, 1.0, self.n_embd)
        ddd = tf.reshape(ddd, (1, 1, self.n_embd))

        # Initialize time-related parameters
        time_maa_x_init = (1.0 - tf.pow(ddd, 0.6 * tf.pow(ratio_1_to_almost0, 0.9))).numpy()
        self.time_maa_x = self.add_weight(
            shape=(1, 1, self.n_embd),
            initializer=tf.constant_initializer(time_maa_x_init),
            trainable=True,
            name='time_maa_x'
        )
        time_maa_rg_init = (1.0 - tf.pow(ddd, 0.2 * ratio_1_to_almost0)).numpy()
        self.time_maa_rg = self.add_weight(
            shape=(1, 1, self.n_embd),
            initializer=tf.constant_initializer(time_maa_rg_init),
            trainable=True,
            name='time_maa_rg'
        )
        time_maa_wa_init = (1.0 - tf.pow(ddd, 0.9 * ratio_1_to_almost0)).numpy()
        self.time_maa_wa = self.add_weight(
            shape=(1, 1, self.n_embd),
            initializer=tf.constant_initializer(time_maa_wa_init),
            trainable=True,
            name='time_maa_wa'
        )
        time_maa_k_init = (1.0 - (tf.pow(ddd, 0.9 * ratio_1_to_almost0) + 0.4 * ratio_0_to_1)).numpy()
        self.time_maa_k = self.add_weight(
            shape=(1, 1, self.n_embd),
            initializer=tf.constant_initializer(time_maa_k_init),
            trainable=True,
            name='time_maa_k'
        )
        time_maa_v_init = (1.0 - (tf.pow(ddd, 0.4 * ratio_1_to_almost0) + 0.6 * ratio_0_to_1)).numpy()
        self.time_maa_v = self.add_weight(
            shape=(1, 1, self.n_embd),
            initializer=tf.constant_initializer(time_maa_v_init),
            trainable=True,
            name='time_maa_v'
        )

        # Initialize decay_speed
        decay_speed = (-6.5 + 5 * tf.pow(tf.linspace(0.0, 1.0, self.dim_att), 0.85 + tf.pow(ratio_0_to_1, 0.5))).numpy()
        self.time_decay = self.add_weight(
            shape=(1, 1, self.dim_att),
            initializer=tf.constant_initializer(decay_speed),
            trainable=True,
            name='time_decay'
        )

        # Initialize other time-related parameters
        self.time_faaaa = self.add_weight(
            shape=(1, 1, self.n_head, self.head_size),
            initializer='zeros',
            trainable=True,
            name='time_faaaa'
        )
        self.time_aaaaa = self.add_weight(
            shape=(1, 1, self.dim_att),
            initializer='zeros',
            trainable=True,
            name='time_aaaaa'
        )

        # Initialize LORA-like projections
        D_MIX_LORA = 32
        self.time_maa_w1 = self.add_weight(
            shape=(self.n_embd, D_MIX_LORA * 4),
            initializer='zeros',
            trainable=True,
            name='time_maa_w1'
        )
        self.time_maa_w2 = self.add_weight(
            shape=(4, D_MIX_LORA, self.n_embd),
            initializer=tf.keras.initializers.Orthogonal(gain=math.sqrt(self.n_embd / (D_MIX_LORA * 4)) * 0.1),
            trainable=True,
            name='time_maa_w2'
        )

        D_DECAY_LORA = 64
        self.time_decay_w1 = self.add_weight(
            shape=(self.n_embd, D_DECAY_LORA),
            initializer='zeros',
            trainable=True,
            name='time_decay_w1'
        )
        self.time_decay_w2 = self.add_weight(
            shape=(D_DECAY_LORA, self.dim_att),
            initializer=tf.keras.initializers.Orthogonal(gain=0.1),
            trainable=True,
            name='time_decay_w2'
        )

        D_AAA_LORA = 16
        self.time_aaa_w1 = self.add_weight(
            shape=(self.n_embd, D_AAA_LORA),
            initializer='zeros',
            trainable=True,
            name='time_aaa_w1'
        )
        self.time_aaa_w2 = self.add_weight(
            shape=(D_AAA_LORA, self.dim_att),
            initializer=tf.keras.initializers.Orthogonal(gain=0.1),
            trainable=True,
            name='time_aaa_w2'
        )

        D_KKK_LORA = 16
        self.time_kkk_w1 = self.add_weight(
            shape=(self.n_embd, D_KKK_LORA),
            initializer='zeros',
            trainable=True,
            name='time_kkk_w1'
        )
        self.time_kkk_w2 = self.add_weight(
            shape=(D_KKK_LORA, self.dim_att),
            initializer=tf.keras.initializers.Orthogonal(gain=0.1),
            trainable=True,
            name='time_kkk_w2'
        )

        D_GATE_LORA = 128
        self.gate_w1 = self.add_weight(
            shape=(self.n_embd, D_GATE_LORA),
            initializer=tf.keras.initializers.Orthogonal(gain=0.1),
            trainable=True,
            name='gate_w1'
        )
        self.gate_w2 = self.add_weight(
            shape=(D_GATE_LORA, self.dim_att),
            initializer=tf.keras.initializers.Orthogonal(gain=0.1),
            trainable=True,
            name='gate_w2'
        )

        D_MA_LORA = 16
        self.ma_w1 = self.add_weight(
            shape=(self.n_embd, D_MA_LORA),
            initializer='zeros',
            trainable=True,
            name='ma_w1'
        )
        self.ma_w2 = self.add_weight(
            shape=(D_MA_LORA, self.dim_att),
            initializer=tf.keras.initializers.Orthogonal(gain=0.1),
            trainable=True,
            name='ma_w2'
        )
        self.time_misc_a = self.add_weight(
            shape=(1, 1, self.n_embd),
            initializer='zeros',
            trainable=True,
            name='time_misc_a'
        )
        D_MK_LORA = 16
        self.mk_w1 = self.add_weight(
            shape=(self.n_embd, D_MK_LORA),
            initializer='zeros',
            trainable=True,
            name='mk_w1'
        )
        self.mk_w2 = self.add_weight(
            shape=(D_MK_LORA, self.dim_att),
            initializer=tf.keras.initializers.Orthogonal(gain=0.1),
            trainable=True,
            name='mk_w2'
        )
        self.time_misc_k = self.add_weight(
            shape=(1, 1, self.n_embd),
            initializer='zeros',
            trainable=True,
            name='time_misc_k'
        )

        # Learnable transformations
        self.receptance = layers.Dense(self.dim_att, use_bias=False, name='receptance')
        self.key = layers.Dense(self.dim_att, use_bias=False, name='key')
        self.value = layers.Dense(self.dim_att, use_bias=False, name='value')
        self.output_dense = layers.Dense(self.n_embd, use_bias=False, name='output')

        # Group Normalization
        self.ln_x = layers.GroupNormalization(epsilon=1e-5, name='ln_x')

        # Build weights so that we can intialise them correctly later on.
        self.receptance.build((None, None, self.n_embd))
        self.key.build((None, None, self.n_embd))
        self.value.build((None, None, self.n_embd))
        self.output_dense.build((None, None, self.dim_att))

        # Initialize weights similar to PyTorch, this could have been done above but it makes the code less readable.
        self.receptance.kernel.assign(tf.random.uniform([self.n_embd, self.dim_att], -0.5 / math.sqrt(self.n_embd), 0.5 / math.sqrt(self.n_embd)))
        self.key.kernel.assign(tf.random.uniform([self.n_embd, self.dim_att], -0.05 / math.sqrt(self.n_embd), 0.05 / math.sqrt(self.n_embd)))
        self.value.kernel.assign(tf.random.uniform([self.n_embd, self.dim_att], -0.5 / math.sqrt(self.n_embd), 0.5 / math.sqrt(self.n_embd)))
        self.output_dense.kernel.assign(tf.zeros([self.dim_att, self.n_embd]))

    def call(self, x, training=False):
        B = tf.shape(x)[0]
        T = tf.shape(x)[1]
        C = tf.shape(x)[2]
        H = self.n_head
        N = self.head_size

        # Time shift
        x_shifted = tf.concat([tf.zeros_like(x[:, :1, :]), x[:, :-1, :]], axis=1) - x
        
        # Apply time_maa_* parameters
        xxx = x + x_shifted * self.time_maa_x  # Broadcasted multiplication
        xxx = tf.tanh(tf.matmul(xxx, self.time_maa_w1))
        xxx = tf.reshape(xxx, [B * T, 4, -1])
        xxx = tf.transpose(xxx, perm=[1, 0, 2])  # Shape: [4, B*T, ...]
        xxx = tf.matmul(xxx, self.time_maa_w2)
        xxx = tf.reshape(xxx, [4, B, T, -1])
        mrg, mwa, mk, mv = tf.unstack(xxx, axis=0)  # Each has shape [B, T, ...]
        # Compute transformed inputs
        xrg = x + x_shifted * (self.time_maa_rg + mrg)
        xwa = x + x_shifted * (self.time_maa_wa + mwa)
        xk = x + x_shifted * (self.time_maa_k + mk)
        xv = x + x_shifted * (self.time_maa_v + mv)

        # Compute receptive parameters
        r = self.receptance(xrg)  # Shape: [B, T, C]
        w_input = self.time_decay + tf.matmul(xwa, self.time_decay_w1) @ self.time_decay_w2
        w = -tf.nn.softplus(-w_input) - 0.5
        k = self.key(xk)  # Shape: [B, T, C]
        v = self.value(xv)  # Shape: [B, T, C]
        g = tf.nn.tanh(tf.matmul(xrg, self.gate_w1)) @ self.gate_w2
        # Normalize keys
        kk = k + tf.matmul(tf.nn.tanh(tf.matmul(xk, self.time_kkk_w1)), self.time_kkk_w2)
        kk = tf.nn.l2_normalize(tf.reshape(kk, [B, T, H, -1]), axis=-1)
        kk = tf.reshape(kk, [B, T, C])

        a = tf.sigmoid(self.time_aaaaa + (tf.matmul(xwa, self.time_aaa_w1) @ self.time_aaa_w2))

        # Mix Keys
        ma = tf.sigmoid(self.time_misc_a + (tf.matmul(xwa, self.ma_w1) @ self.ma_w2))
        k = k * ma + k * a * (1 - ma)
        mk = tf.sigmoid(self.time_misc_k + (tf.matmul(xk, self.mk_w1) @ self.mk_w2))
        k = k * tf.clip_by_value(tf.exp(-w * mk), 0, tf.float32.max)  # Similar to clamp in PyTorch

        # Initialize state
        initial_state = tf.zeros((B, H, N), dtype=x.dtype)

        # Reshape to (B, T, H, N)
        r_reshaped = tf.reshape(r, (B, T, H, N))
        w_reshaped = tf.reshape(w, (B, T, H, N))
        k_reshaped = tf.reshape(k, (B, T, H, N))
        v_reshaped = tf.reshape(v, (B, T, H, N))
        a_reshaped = tf.reshape(-kk, (B, T, H, N))
        b_reshaped = tf.reshape(kk*a, (B, T, H, N))

        # Define the scan function
        def scan_fn(prev_state, inputs):
            a_t, w_t, k_t, v_t, b_t, r_t = inputs
            sa = tf.reduce_sum(a_t * prev_state, axis=-1, keepdims=True)  # (B, H, 1)

            new_state = prev_state * w_t + k_t * v_t + sa * b_t  # (B, H, N)

            y_t = new_state * r_t  # (B, H, N)
            return y_t

        a_t = tf.transpose(a_reshaped, perm=[1, 0, 2, 3])  # (T, B, H, N)
        w_t = tf.transpose(w_reshaped, perm=[1, 0, 2, 3])  # (T, B, H, N)
        k_t = tf.transpose(k_reshaped, perm=[1, 0, 2, 3])  # (T, B, H, N)
        v_t = tf.transpose(v_reshaped, perm=[1, 0, 2, 3])  # (T, B, H, N)
        b_t = tf.transpose(b_reshaped, perm=[1, 0, 2, 3])  # (T, B, H, N)
        r_t = tf.transpose(r_reshaped, perm=[1, 0, 2, 3])  # (T, B, H, N)

        # Stack inputs
        scan_inputs = (a_t, tf.exp(-tf.exp(w_t)), k_t, v_t, b_t, r_t)

        # Perform scan
        outputs = tf.scan(
          scan_fn,
          elems=scan_inputs,
          initializer=initial_state,
          swap_memory=True
        ) #(T, B, H, N)
        
        outputs = tf.transpose(outputs, perm=[1, 0, 2, 3])  # Shape: [B, T, H, N]

        # Apply Group Normalization
        outputs = self.ln_x(outputs)
        tf.print(outputs)
        outputs = tf.reshape(outputs + tf.reduce_sum(r_reshaped * k_reshaped * self.time_faaaa, axis=-1, keepdims=True) * v_reshaped, shape=(B, T, C))

        # Compute final output
        x = self.output_dense(outputs * g)
        return x


class MLP(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
        super(MLP, self).__init__(**kwargs)
        self.c_fc = layers.Dense(7 * config.n_embd // 2, use_bias=False, name='c_fc')
        self.c_proj = layers.Dense(config.n_embd, use_bias=False, name='c_proj', kernel_initializer='zeros')
    def call(self, x):
        x = self.c_fc(x)
        x = tf.nn.relu(x) ** 2
        x = self.c_proj(x)
        return x

class Block(tf.keras.layers.Layer):
    def __init__(self, config, layer_id, **kwargs):
        super(Block, self).__init__(**kwargs)
        self.attn = RWKV7(config, layer_id=layer_id, name='attn')
        self.mlp = MLP(config, name='mlp')
        self.ln1 = layers.LayerNormalization(epsilon=1e-5, name='ln1')
        self.ln2 = layers.LayerNormalization(epsilon=1e-5, name='ln2')

    def call(self, x, training=False):
        # Attention with residual connection
        attn_output = self.attn(self.ln1(x), training=training)
        x = x + attn_output

        # MLP with residual connection
        mlp_output = self.mlp(self.ln2(x), training=training)
        x = x + mlp_output
        return x