import tensorflow as tf import numpy as np class SquaredReLU(tf.keras.layers.Layer): def __init__(self, **kwargs): super(SquaredReLU, self).__init__(**kwargs) def call(self, x): return tf.square(tf.nn.relu(x)) class TokenShift(tf.keras.layers.Layer): """ Position encoding via Token Shifting. Shifts the sequence to provide local positional context without additive embeddings. """ def __init__(self, shift_amount=1, **kwargs): super(TokenShift, self).__init__(**kwargs) self.shift_amount = shift_amount def call(self, x): # x shape: [batch, seq_len, dim] seq_len = tf.shape(x)[1] shifted = tf.roll(x, shift=self.shift_amount, axis=1) # Create a mask to zero out the elements that wrapped around # The first 'shift_amount' elements are the ones that wrapped from the end mask = tf.ones([seq_len], dtype=x.dtype) mask = tf.concat([tf.zeros([self.shift_amount], dtype=x.dtype), mask[self.shift_amount:]], axis=0) mask = tf.reshape(mask, [1, seq_len, 1]) return x + (shifted * mask) class GroupNorm(tf.keras.layers.Layer): """ Group Normalization instead of LayerNorm. """ def __init__(self, groups=32, eps=1e-6, **kwargs): super(GroupNorm, self).__init__(**kwargs) self.groups = groups self.eps = eps def build(self, input_shape): dim = input_shape[-1] if dim % self.groups != 0: raise ValueError(f"Dimension {dim} must be divisible by groups {self.groups}") self.gamma = self.add_weight(shape=(dim,), initializer='ones', trainable=True, name='gamma') self.beta = self.add_weight(shape=(dim,), initializer='zeros', trainable=True, name='beta') def call(self, x): # x shape: [batch, seq_len, dim] batch, seq_len, dim = tf.shape(x)[0], tf.shape(x)[1], tf.shape(x)[2] x = tf.reshape(x, [batch, seq_len, self.groups, dim // self.groups]) mean = tf.reduce_mean(x, axis=-1, keepdims=True) var = tf.reduce_mean(tf.square(x - mean), axis=-1, keepdims=True) x = (x - mean) / tf.sqrt(var + self.eps) x = tf.reshape(x, [batch, seq_len, dim]) return x * self.gamma + self.beta class StochasticDepth(tf.keras.layers.Layer): """ Regularization: Randomly drops layers during training. """ def __init__(self, drop_prob=0.1, **kwargs): super(StochasticDepth, self).__init__(**kwargs) self.drop_prob = drop_prob def call(self, x, training=False): if not training or self.drop_prob == 0: return x keep_prob = 1.0 - self.drop_prob random_tensor = tf.random.uniform(shape=[tf.shape(x)[0]], minval=0, maxval=1) binary_mask = tf.cast(random_tensor < keep_prob, tf.float32) # Reshape mask for broadcasting: [batch, 1, 1] mask = tf.reshape(binary_mask, [-1, 1, 1]) return x * mask / keep_prob class TimeMix(tf.keras.layers.Layer): """ Unique sequence mixing mechanism for Tera v3. Uses a depth-wise temporal convolution to mix information across time. """ def __init__(self, dim, **kwargs): super(TimeMix, self).__init__(**kwargs) self.dim = dim # Depth-wise convolution for temporal mixing self.conv = tf.keras.layers.DepthwiseConv1D( kernel_size=3, padding='same', activation=None ) self.norm = GroupNorm() self.proj = tf.keras.layers.Dense(dim) def call(self, x): # x shape: [batch, seq_len, dim] residual = x x = self.norm(x) x = self.conv(x) x = self.proj(x) return x + residual class ChannelMix(tf.keras.layers.Layer): """ Feed-forward network using Squared ReLU. """ def __init__(self, dim, expand_factor=4, **kwargs): super(ChannelMix, self).__init__(**kwargs) self.dim = dim self.inner_dim = dim * expand_factor self.norm = GroupNorm() self.fc1 = tf.keras.layers.Dense(self.inner_dim) self.sq_relu = SquaredReLU() self.fc2 = tf.keras.layers.Dense(dim) def call(self, x): residual = x x = self.norm(x) x = self.fc1(x) x = self.sq_relu(x) x = self.fc2(x) return x + residual class TeraBlock(tf.keras.layers.Layer): """ The core block of Tera v3. """ def __init__(self, dim, drop_prob=0.1, **kwargs): super(TeraBlock, self).__init__(**kwargs) self.token_shift = TokenShift() self.time_mix = TimeMix(dim) self.channel_mix = ChannelMix(dim) self.stoch_depth = StochasticDepth(drop_prob) def call(self, x, training=False): # Positional encoding via shift x = self.token_shift(x) # Sequence mixing x = x + self.stoch_depth(self.time_mix(x), training=training) # Feed-forward mixing x = x + self.stoch_depth(self.channel_mix(x), training=training) return x class TeraVisionEncoder(tf.keras.layers.Layer): """ Unique Vision Capability: Spectral-Spatial Shift Projection. Instead of standard ViT, it uses a custom patch-shifting mechanism. """ def __init__(self, patch_size=16, embed_dim=512, **kwargs): super(TeraVisionEncoder, self).__init__(**kwargs) self.patch_size = patch_size self.embed_dim = embed_dim # Use a custom convolution for patch embedding to avoid 'Google' style ViT layers self.patch_embed = tf.keras.layers.Conv2D( filters=embed_dim, kernel_size=patch_size, strides=patch_size, padding='valid' ) def call(self, images): # images: [batch, h, w, c] x = self.patch_embed(images) # [batch, h/p, w/p, dim] # Flatten spatial dimensions to sequence batch = tf.shape(x)[0] x = tf.reshape(x, [batch, -1, self.embed_dim]) return x class TeraV3(tf.keras.Model): """ Tera v3 Language Model. """ def __init__(self, vocab_size, dim=512, depth=12, embed_dim=512, **kwargs): super(TeraV3, self).__init__(**kwargs) # Untied Embeddings: Input and Output embeddings are separate self.input_embedding = tf.keras.layers.Embedding(vocab_size, embed_dim) self.output_projection = tf.keras.layers.Dense(vocab_size) # Vision Encoder self.vision_encoder = TeraVisionEncoder(embed_dim=dim) # Core Architecture self.blocks = [TeraBlock(dim) for _ in range(depth)] self.final_norm = GroupNorm() # Project input embedding to model dim if they differ if embed_dim != dim: self.input_proj = tf.keras.layers.Dense(dim) else: self.input_proj = tf.keras.layers.Lambda(lambda x: x) def call(self, inputs, training=False, vision_inputs=None): # Handle vision inputs if vision_inputs is not None: v_emb = self.vision_encoder(vision_inputs) # Concat vision tokens with text tokens x = self.input_embedding(inputs) x = self.input_proj(x) x = tf.concat([v_emb, x], axis=1) else: x = self.input_embedding(inputs) x = self.input_proj(x) # Pass through Tera Blocks for block in self.blocks: x = block(x, training=training) x = self.final_norm(x) # Project to vocab for text generation # For vision-text models, usually only the text part produces logits # but for simplicity, we project the whole sequence. logits = self.output_projection(x) return logits # Example instantiation if __name__ == "__main__": vocab_size = 10000 model = TeraV3(vocab_size=vocab_size) # Dummy text input text_in = tf.constant([[1, 2, 3, 4, 5]], dtype=tf.int32) # Dummy vision input vision_in = tf.random.uniform([1, 224, 224, 3]) output = model(text_in, vision_inputs=vision_in) print("Output shape:", output.shape) # [batch, (num_patches + seq_len), vocab_size]