Rename MiniLLM.py to NanoLLM.py
Browse files- MiniLLM.py → NanoLLM.py +1 -1
MiniLLM.py → NanoLLM.py
RENAMED
|
@@ -34,7 +34,7 @@ class TokenAndPositionEmbedding(nnx.Module):
|
|
| 34 |
positions = jnp.arange(seq_len)[None, :]
|
| 35 |
return self.token_emb(x) + self.pos_emb(positions)
|
| 36 |
|
| 37 |
-
class
|
| 38 |
|
| 39 |
def __init__(self, maxlen=maxlen, vocab_size=vocab_size, embed_dim=embed_dim, num_heads=num_heads,
|
| 40 |
feed_forward_dim=feed_forward_dim, num_transformer_blocks=num_transformer_blocks, *, rngs=nnx.Rngs(0)):
|
|
|
|
| 34 |
positions = jnp.arange(seq_len)[None, :]
|
| 35 |
return self.token_emb(x) + self.pos_emb(positions)
|
| 36 |
|
| 37 |
+
class NanoLLM(nnx.Module):
|
| 38 |
|
| 39 |
def __init__(self, maxlen=maxlen, vocab_size=vocab_size, embed_dim=embed_dim, num_heads=num_heads,
|
| 40 |
feed_forward_dim=feed_forward_dim, num_transformer_blocks=num_transformer_blocks, *, rngs=nnx.Rngs(0)):
|