Help optimise the attention head part:
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.K = 8
        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), tf.float32.min, 0)  # Similar to clamp in PyTorch

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

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

        # Prepare inputs for tf.scan, transposing to [T//K, B, K, H, N]
        scan_inputs = (
            tf.transpose(w_reshaped, perm=[1, 0, 2, 3, 4]),  # w
            tf.transpose(r_reshaped, perm=[1, 0, 2, 3, 4]),  # r (q in ref code)
            tf.transpose(k_reshaped, perm=[1, 0, 2, 3, 4]),  # k
            tf.transpose(v_reshaped, perm=[1, 0, 2, 3, 4]),  # v
            tf.transpose(a_reshaped, perm=[1, 0, 2, 3, 4]),  # a
            tf.transpose(b_reshaped, perm=[1, 0, 2, 3, 4])   # b
        )


        def scan_fn(state_and_output, inputs_t):
            state, _ = state_and_output #unpack state and discard carried output.
            sw, sq, sk, sv, sa, sb = inputs_t

            w_chunk = tf.exp(-tf.exp(sw))
            fw = tf.reduce_prod(w_chunk, axis=2, keepdims=True)
            incl_pref = tf.math.cumprod(w_chunk, axis=2, exclusive=False)
            non_incl_pref = incl_pref / (w_chunk + 1e-6)
            inv_incl_pref = 1 / (incl_pref + 1e-6)

            wq = sq * incl_pref
            wa = sa * non_incl_pref
            kwi = sk * inv_incl_pref
            bwi = sb * inv_incl_pref

            wa_bwiT = tf.matmul(wa, bwi, transpose_b=True)
            ab = tf.linalg.band_part(wa_bwiT, -1, 0)
            ak = tf.linalg.band_part(tf.matmul(wa, kwi, transpose_b=True), -1, 0)

            K_actual = tf.shape(sw)[2]
            identity = tf.eye(K_actual, batch_shape=tf.shape(sw)[:2], dtype=wa_bwiT.dtype)
            ab_inv = tf.linalg.triangular_solve(-ab, identity, lower=True)
            ab_u = tf.matmul(ak, sv) + tf.matmul(wa, state, transpose_b=True)
            u = tf.matmul(ab_inv, ab_u)

            qk = tf.linalg.band_part(tf.matmul(wq, kwi, transpose_b=True), 0, -1)
            qb = tf.linalg.band_part(tf.matmul(wq, bwi, transpose_b=True), 0, -1)
            yy = tf.matmul(qk, sv) + tf.matmul(qb, u) + tf.matmul(wq, state)

            new_state = state * fw + tf.matmul(sv, kwi * fw, transpose_a=True) + tf.matmul(u, bwi * fw, transpose_a=True)

            return (new_state, yy)


        # Initialize state for tf.scan
        initial_state = tf.zeros((B, self.K, N, N), dtype=x.dtype)  # Initial state shape (B, H, N)

        _, y = tf.scan(
            scan_fn,
            elems=scan_inputs,
            initializer=(initial_state, tf.zeros((B,self.K, H, N), dtype=x.dtype)), #Dummy initializer for output
        )
        
        outputs = tf.reshape(tf.transpose(y, perm=[1,0,2,3,4]), (B * T,C))  # Shape: [B, T, H, N]

        # Apply Group Normalization
        outputs = tf.reshape(self.ln_x(outputs), (B, T, C))
        outputs = outputs +  tf.reshape(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

class GPT(tf.keras.Model):
    def __init__(self, config, **kwargs):
        super(GPT, self).__init__(**kwargs)
        self.config = config
        self.transformer_wte = layers.Embedding(config.vocab_size, config.n_embd, name='wte', mask_zero=True)
        self.transformer_h = [Block(config, layer_id=i, name=f'h_{i}') for i in range(config.n_layer)]
        self.ln_out = layers.LayerNormalization(epsilon=1e-5, name='ln_out')
        self.lm_head = layers.Dense(config.vocab_size, use_bias=False, name='lm_head')
        self.loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    def call(self, idx, training=False):
        # Forward pass through embeddings
        x = self.transformer_wte(idx)  # Shape: [B, T, C]
        # Forward pass through transformer blocks
        for block in self.transformer_h:
            x = block(x, training=training)
       
        # Final layer normalization
        x = self.ln_out(x)
        logits = self.lm_head(x)
        #logits = tf.cast(logits, tf.float32)  # Ensure logits are in float32
        return logits[:, :-1, :]

@dataclass
class GPTConfig:
    vocab_size : int = 50304
    n_layer : int = 12
    n_head : int = 768//HEAD_SIZE
    n_embd : int = 768