import tensorflow as tf
from tensorflow.keras import layers, Model

class Config:
    def __init__(self, n_embed=50304, d_model=768, n_heads=12, n_layers=12):
        self.n_embed = n_embed
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_layers = n_layers

def lerp(start, end, weight):
    return start + weight * (end - start)

class LoRA_MLP(layers.Layer):
    def __init__(self, dim, dim_hidden, init_value=None):
        super(LoRA_MLP, self).__init__()
        if init_value is None:
            init_value = tf.zeros(dim)
        self.base = tf.Variable(init_value, trainable=True)
        self.W_a = layers.Dense(dim_hidden, use_bias=False)
        self.W_b = layers.Dense(dim, use_bias=False)

    def call(self, x):
        return tf.cast(self.base, x.dtype) + self.W_b(tf.nn.tanh(self.W_a(x)))

class DDLerp(layers.Layer):
    def __init__(self, dim, dim_hidden):
        super(DDLerp, self).__init__()
        self.lora = LoRA_MLP(dim, dim_hidden)

    def call(self, x_premixed, x, x_shifted_one_to_the_past):
        y = self.lora(x_premixed)
        y = lerp(x, x_shifted_one_to_the_past, y)
        return y

class RWKV(Model):
    def __init__(self, cfg=Config()):
        super(RWKV, self).__init__()
        self.cfg = cfg
        self.loss_fn = tf.losses.SparseCategoricalCrossentropy(from_logits=True)
        self.embed = tf.Variable(tf.ones((cfg.n_embed, cfg.d_model)) * 1e-4, trainable=True)
        self.embed_norm = layers.LayerNormalization()
        self.blocks = [Block(cfg, layer_id) for layer_id in range(cfg.n_layers)]
        self.out_norm = layers.LayerNormalization()
        self.unembed = layers.Dense(cfg.n_embed, use_bias=False)

    def call(self, x, s=None):
        x = tf.nn.embedding_lookup(self.embed, x)
        x = self.embed_norm(x)

        for block in self.blocks:
            x = block(x)

        x = self.out_norm(x)
        x = self.unembed(x)

        return x[:, :-1, :]

class Block(layers.Layer):
    def __init__(self, cfg, layer_id):
        super(Block, self).__init__()
        self.time_mixer = TimeMixer(cfg, layer_id)
        self.channel_mixer = ChannelMixer(cfg, layer_id)
        self.C = cfg.d_model
        self.H = cfg.n_heads
        self.K = cfg.d_model // cfg.n_heads

    def call(self, x):
        B = tf.shape(x)[0]
        time_mixer_x_state = tf.zeros((B, self.C), dtype=x.dtype)
        vk_state = tf.zeros((B, self.H, self.K, self.K), dtype=x.dtype)
        channel_mixer_x_state = tf.zeros((B, self.C), dtype=x.dtype)

        x = self.time_mixer(x, time_mixer_x_state, vk_state)
        x = self.channel_mixer(x, channel_mixer_x_state)
        return x

class LoRA_Simple(layers.Layer):
    def __init__(self, dim, dim_hidden, init_value=None):
        super(LoRA_Simple, self).__init__()
        if init_value is None:
            init_value = tf.zeros(dim)
        self.base = tf.Variable(init_value, trainable=True)
        self.W_a = layers.Dense(dim_hidden, use_bias=False)
        self.W_b = layers.Dense(dim, use_bias=False)

    def call(self, x):
        return tf.cast(self.base, x.dtype) + self.W_b(self.W_a(x))

class TimeMixer(layers.Layer):
    def __init__(self, cfg, layer_id):
        super(TimeMixer, self).__init__()
        self.cfg = cfg
        self.layer_id = layer_id

        d_attn = d_model = cfg.d_model

        self.time_mixer_prenorm = layers.LayerNormalization()

        self.W_ddlerp_premix = self.add_weight(
            shape=(d_model,), initializer='zeros', trainable=True
        )

        self.ddlerps = [DDLerp(d_model, 32) for _ in range(6)]
        self.decay_lora = LoRA_MLP(d_model, 64, tf.ones(d_model))
        self.W_proj_r = layers.Dense(d_model, use_bias=False)
        self.W_proj_k = layers.Dense(d_model, use_bias=False)
        self.W_proj_v = layers.Dense(d_model, use_bias=False)
        self.W_proj_out = layers.Dense(d_model, use_bias=False)

        self.decay_lora_mlp = LoRA_MLP(
            d_model, 64 if d_model < 4096 else 128
        )

        self.iclr_lora = LoRA_MLP(d_model, 64)

        D_DEFORMED_KEY_LORA = 64
        self.deformed_key_w1 = self.add_weight(
            shape=(d_model, D_DEFORMED_KEY_LORA),
            initializer='zeros',
            trainable=True
        )
        self.deformed_key_w2 = self.add_weight(
            shape=(D_DEFORMED_KEY_LORA, d_attn),
            initializer=tf.random_uniform_initializer(-0.01, 0.01),
            trainable=True
        )

        D_GATE_LORA = 128
        self.gate_w1 = self.add_weight(
            shape=(d_model, D_GATE_LORA),
            initializer='zeros',
            trainable=True
        )
        self.gate_w2 = self.add_weight(
            shape=(D_GATE_LORA, d_attn),
            initializer=tf.random_uniform_initializer(-0.01, 0.01),
            trainable=True
        )

        self.iclr_mix_amt_lora = LoRA_MLP(d_model, 16)
        self.one_minus_decay_mix_amt_lora = LoRA_MLP(d_model, 16)

        self.u = self.add_weight(
            shape=(cfg.n_heads, cfg.d_model // cfg.n_heads),
            initializer='ones',
            trainable=True
        )

        self.group_norm = layers.LayerNormalization(epsilon=64e-5)

    def call(self, x, x_state, vk_state):
        x = self.time_mixer_prenorm(x)
        B, T = tf.shape(x)[0], tf.shape(x)[1]
        C = self.cfg.d_model
        H = self.cfg.n_heads
        K = C // H

        # Shift x to get x_shifted_one_to_the_past
        x_shifted_one_to_the_past = tf.concat(
            [tf.expand_dims(x_state, 1), x[:, :-1, :]], axis=1
        )

        # Compute x_premixed
        W_ddlerp_premix_cast = tf.cast(self.W_ddlerp_premix, dtype=x.dtype)
        x_premixed = lerp(x, x_shifted_one_to_the_past, W_ddlerp_premix_cast)

        # Compute ddlerp outputs
        ddlerp_outputs = [
            ddlerp(x_premixed, x, x_shifted_one_to_the_past)
            for ddlerp in self.ddlerps
        ]
        x_receptance, x_key, x_value, x_decay, x_iclr, x_gate = ddlerp_outputs

        # Compute projections
        r = self.W_proj_r(x_receptance)
        k = self.W_proj_k(x_key)
        v = self.W_proj_v(x_value)

        # Compute gate
        gate = tf.tanh(tf.matmul(x_gate, self.gate_w1))
        gate = tf.matmul(gate, self.gate_w2)

        # Compute decay
        log_neglog_of_decay = self.decay_lora_mlp(x_decay)
        log_neglog_of_decay = -0.5 - tf.nn.softplus(-log_neglog_of_decay)
        log_of_decay = -tf.exp(log_neglog_of_decay)
        decay = tf.exp(log_of_decay)
        decay = tf.reshape(decay, [B, T, H, -1])

        # Compute iclr
        iclr = tf.sigmoid(self.iclr_lora(x_iclr))

        deformed_key = k + tf.matmul(tf.tanh(tf.matmul(x_key, self.deformed_key_w1)), self.deformed_key_w2)
        deformed_key = tf.reshape(deformed_key, [B, T, H, -1])
        deformed_key = tf.nn.l2_normalize(deformed_key, axis=-1)

        iclr_mix_amt = tf.sigmoid(self.iclr_mix_amt_lora(x_iclr))
        k = lerp(k, k * iclr, iclr_mix_amt)

        one_minus_decay_mix_amt = tf.sigmoid(self.one_minus_decay_mix_amt_lora(x_key))
        k = k * tf.exp(tf.minimum(log_of_decay * one_minus_decay_mix_amt, 0))

        # Reshape into heads
        r = tf.reshape(r, [B, T, H, -1])
        k = tf.reshape(k, [B, T, H, -1])
        v = tf.reshape(v, [B, T, H, -1])
        decay = tf.reshape(decay, [B, T, H, -1, 1])
        iclr = tf.reshape(iclr, [B, T, H, -1, 1])
        deformed_key = tf.reshape(deformed_key, [B, T, H, -1, 1])

        r_t = tf.transpose(r, [1, 0, 2, 3])
        k_t = tf.transpose(k, [1, 0, 2, 3])
        v_t = tf.transpose(v, [1, 0, 2, 3])
        # Prepare per-timestep inputs
        per_timestep_inputs = {
            'decay_t': tf.transpose(decay, [1, 0, 2, 3, 4]),
            'iclr_t': tf.transpose(iclr, [1, 0, 2, 3, 4]),
            'deformed_key_t': tf.transpose(deformed_key, [1, 0, 2, 3, 4]),
            'kv_t': tf.matmul(tf.expand_dims(v_t, -1), tf.expand_dims(k_t, -1), transpose_a=True),
        }

        vk_state_0 = vk_state # Initial state

        def time_mixer_step(vk_state_prev, per_timestep_input):
            decay_t = per_timestep_input['decay_t']
            iclr_t = per_timestep_input['iclr_t']
            deformed_key_t = per_timestep_input['deformed_key_t']
            kv_t = per_timestep_input['kv_t']

            vk_state_new = vk_state_prev * decay_t
            vk_state_new -= tf.matmul(
                tf.matmul(vk_state_new, deformed_key_t),
                iclr_t * deformed_key_t,
                transpose_b=True
            )
            vk_state_new += kv_t

            return vk_state_new

        # Run tf.scan
        out = tf.scan(
            fn=time_mixer_step,
            elems=per_timestep_inputs,
            initializer=vk_state_0,
            parallel_iterations=1
        )

        out = tf.transpose(out, [1,0,2,3,4])

        out = tf.matmul(out, tf.expand_dims(r, -1))
        out = tf.squeeze(out, -1)

        out = tf.reshape(out, [B*T, C])
        out = self.group_norm(out)

        out = tf.reshape(out, [B, T, C])

        bonus = tf.reduce_sum(r_t * k_t * self.u, axis=-1, keepdims=True) * v_t
        bonus = tf.reshape(bonus, [B, T, C])

        out += bonus
        out *= gate

        out = self.W_proj_out(out)

        return x + out

class ChannelMixer(layers.Layer):
    def __init__(self, cfg, layer_id):
        super(ChannelMixer, self).__init__()
        self.cfg = cfg
        self.prenorm = layers.LayerNormalization()
        self.W_in = layers.Dense(cfg.d_model * 4, use_bias=False)
        self.W_out = layers.Dense(cfg.d_model, use_bias=False)
        self.W_lerp_in = layers.Dense(cfg.d_model, use_bias=False)

    def call(self, x, x_state):
        x = self.prenorm(x)
        x_shifted_one_to_the_past = tf.concat([tf.expand_dims(x_state, 1), x[:, :-1, :]], axis=1)
        inx = lerp(x, x_shifted_one_to_the_past, self.W_lerp_in(x))

        hidden = self.W_in(inx)
        hidden = tf.square(tf.nn.relu(hidden))
        out = self.W_out(hidden)

        return x + out




Edit the above code to improve memory usage and performance.