From ea0490972a3f0a8b3d6ccad4ecbdd31173051f7e Mon Sep 17 00:00:00 2001 From: Sam Anklesaria Date: Tue, 26 May 2026 08:30:24 -0500 Subject: [PATCH 1/2] Add RoPE embeddings --- .../api_reference/flax.nnx/nn/attention.rst | 7 +- flax/nnx/__init__.py | 2 + flax/nnx/nn/attention.py | 132 +++++++++++++++++- tests/nnx/nn/attention_test.py | 38 ++++- 4 files changed, 175 insertions(+), 4 deletions(-) diff --git a/docs_nnx/api_reference/flax.nnx/nn/attention.rst b/docs_nnx/api_reference/flax.nnx/nn/attention.rst index 0317f60e9..1f60d70ea 100644 --- a/docs_nnx/api_reference/flax.nnx/nn/attention.rst +++ b/docs_nnx/api_reference/flax.nnx/nn/attention.rst @@ -8,7 +8,12 @@ Attention :module: flax.nnx :class: MultiHeadAttention +.. flax_module:: + :module: flax.nnx + :class: RoPE + .. autofunction:: combine_masks .. autofunction:: dot_product_attention +.. autofunction:: dot_product_attention_with_rope .. autofunction:: make_attention_mask -.. autofunction:: make_causal_mask \ No newline at end of file +.. autofunction:: make_causal_mask diff --git a/flax/nnx/__init__.py b/flax/nnx/__init__.py index 78a550411..fd02e8e39 100644 --- a/flax/nnx/__init__.py +++ b/flax/nnx/__init__.py @@ -115,6 +115,8 @@ from .nn.attention import MultiHeadAttention as MultiHeadAttention from .nn.attention import combine_masks as combine_masks from .nn.attention import dot_product_attention as dot_product_attention +from .nn.attention import dot_product_attention_with_rope as dot_product_attention_with_rope +from .nn.attention import RoPE as RoPE from .nn.attention import make_attention_mask as make_attention_mask from .nn.attention import make_causal_mask as make_causal_mask from .nn.recurrent import RNNCellBase as RNNCellBase diff --git a/flax/nnx/nn/attention.py b/flax/nnx/nn/attention.py index 4f0c4f0cd..01dbe56a5 100644 --- a/flax/nnx/nn/attention.py +++ b/flax/nnx/nn/attention.py @@ -319,6 +319,128 @@ def reshape_4d(x): return out +class RoPE(Module): + """Rotary Position Embedding (RoPE). + + Precomputes and stores cosine and sine frequency tensors for rotary + positional encoding as described in "RoFormer: Enhanced Transformer + with Rotary Position Embedding" (https://arxiv.org/abs/2104.09864). + + Example usage:: + + >>> from flax import nnx + >>> import functools + >>> rope = nnx.RoPE(embedding_size=64, max_seq_len=128) + >>> layer = nnx.MultiHeadAttention( + ... num_heads=8, in_features=512, qkv_features=512, + ... attention_fn=functools.partial( + ... nnx.dot_product_attention_with_rope, rope=rope), + ... decode=False, rngs=nnx.Rngs(0)) + + Args: + theta: Base frequency for sinusoidal functions. Default is 10000. + embedding_size: Optional. Size of each embedding vector (i.e. head_dim). Must be + even. If provided together with ``max_seq_len``, sin/cos tables are + precomputed and cached as module Variables. + max_seq_len: Optional. Maximum sequence length for precomputed frequencies. Must + be provided together with ``embedding_size``. + """ + + def __init__( + self, + theta: float = 10000.0, + embedding_size: int | None = None, + max_seq_len: int | None = None, + ): + self.theta = theta + if sum([embedding_size is None, max_seq_len is None]) % 2 != 0: + raise ValueError('Either both `embedding_size` and `max_seq_len` or none of them must be provided.') + if embedding_size is not None and max_seq_len is not None: + positions = jnp.arange(float(max_seq_len)) + sin, cos = self._rotation_for(embedding_size, positions) + self.cached_sin = nnx.Variable(sin) + self.cached_cos = nnx.Variable(cos) + + @staticmethod + def _rotate_half(x: Array) -> Array: + d_2 = x.shape[-1] // 2 + return jnp.concatenate([-x[..., d_2:], x[..., :d_2]], axis=-1) + + def _rotation_for(self, embedding_size: int, input_positions: Array): + seq_len = input_positions.shape[-1] + if embedding_size <= 0 or embedding_size % 2 != 0: + raise ValueError('`embedding_size` must be positive and even.') + freqs = 1.0 / ( + self.theta ** (jnp.arange(0.0, embedding_size, 2) / embedding_size) + ) + freqs_outer = jnp.outer(input_positions, freqs) + return jnp.sin(freqs_outer), jnp.cos(freqs_outer) + + def __call__(self, x: Array, input_positions: Array | None = None) -> Array: + """Apply rotary positional encoding. + + Args: + x: Input array of shape ``(..., seq_length, embedding_size)``. + input_positions: Optional position array of shape ``(seq_len,)``. + If ``None`` and precomputed tables exist (i.e. ``embedding_size`` + and ``max_seq_len`` were given at init), the cached tables are + used. Otherwise positions default to ``jnp.arange(seq_len)``. + + Returns: + Array of same shape with RoPE applied. + """ + seq_len, embedding_size = x.shape[-2:] + if input_positions is None and hasattr(self, 'cached_sin'): + sin = self.cached_sin.value[:seq_len] + cos = self.cached_cos.value[:seq_len] + else: + if input_positions is None: + input_positions = jnp.arange(float(seq_len)) + elif input_positions.shape[-1] != seq_len: + raise ValueError('`input_positions` must have the same seq_len as `x`.') + sin, cos = self._rotation_for(embedding_size, input_positions) + return x * jnp.tile(cos, (1, 2)) + self._rotate_half(x) * jnp.tile(sin, (1, 2)) + + +def dot_product_attention_with_rope( + query: Array, + key: Array, + value: Array, + *, + rope: RoPE, + input_positions: Array | None = None, + **kwargs, +) -> Array: + """Dot-product attention with Rotary Position Embedding applied to query and key. + + This function has the same signature as :func:`dot_product_attention` + (plus a ``rope`` keyword argument) and can be used as the ``attention_fn`` + of :class:`MultiHeadAttention` via ``functools.partial``:: + + attention_fn=functools.partial(dot_product_attention_with_rope, rope=rope) + + Args: + query: queries of shape ``[batch..., q_length, num_heads, head_dim]``. + key: keys of shape ``[batch..., kv_length, num_heads, head_dim]``. + value: values of shape ``[batch..., kv_length, num_heads, v_dim]``. + rope: A :class:`RoPE` module instance. + **kwargs: Additional keyword arguments forwarded to + :func:`dot_product_attention`. + + Returns: + Output of shape ``[batch..., q_length, num_heads, v_dim]``. + """ + # query/key: [batch..., seq_length, num_heads, head_dim] + # RoPE.__call__ expects (..., seq_length, head_dim), so vmap over heads. + if input_positions is None: + apply = jax.vmap(rope, in_axes=-2, out_axes=-2) + else: + apply = jax.vmap(rope, in_axes=(-2, -1), out_axes=(-2, -1)) + query = apply(query, input_positions) + key = apply(key, input_positions) + return dot_product_attention(query, key, value, **kwargs) + + class MultiHeadAttention(Module): """Multi-head attention. @@ -588,6 +710,7 @@ def __call__( is_causal=False, out_sharding = None, qkv_sharding = None, + input_positions: Array | None = None ): """Applies multi-head dot product attention on the input data. @@ -624,6 +747,9 @@ def __call__( the output linear layer for the output arrays. qkv_sharding: Optional sharding specification to pass to the QKV linear layers for the output arrays. + input_positions: Optional position indices of shape + ``[batch_sizes..., length]`` forwarded to the ``attention_fn``. + Used by :func:`dot_product_attention_with_rope` for RoPE. Returns: output of shape `[batch_sizes..., length, features]`. @@ -736,6 +862,9 @@ def __call__( dropout_rng = None # apply attention + attn_kwargs = {} + if input_positions is not None: + attn_kwargs["input_positions"] = input_positions x = self.attention_fn( query, key, @@ -748,7 +877,8 @@ def __call__( dtype=self.dtype, precision=self.precision, module=self if sow_weights else None, - is_causal=is_causal + is_causal=is_causal, + **attn_kwargs ) # back to the original inputs dimensions out = self.out(x, out_sharding=out_sharding) diff --git a/tests/nnx/nn/attention_test.py b/tests/nnx/nn/attention_test.py index 167fbf8cf..82aa52344 100644 --- a/tests/nnx/nn/attention_test.py +++ b/tests/nnx/nn/attention_test.py @@ -383,7 +383,41 @@ def _run(m): np.testing.assert_allclose(nnx_out, jax_out, atol=1e-3, rtol=1e-3) -if __name__ == '__main__': - absltest.main() +class TestRoPE(absltest.TestCase): + + def test_relative_position_invariance(self): + """Dot product of RoPE-rotated vectors depends only on relative position. + + For any offset d, should equal + . + """ + head_dim = 64 + max_len = 128 + rope = nnx.RoPE() + + k1, k2 = jax.random.split(jax.random.key(7)) + q_vec = jax.random.normal(k1, (head_dim,)) + k_vec = jax.random.normal(k2, (head_dim,)) + + # Place q at position i and k at position j, then shift both by d. + i, j = 5, 12 + d = 37 + def rope_at(vec, pos): + # Build a dummy sequence long enough, apply RoPE, extract one position. + seq = jnp.zeros((pos + 1, head_dim)).at[pos].set(vec) + return rope(seq)[pos] + q_rot = rope_at(q_vec, i) + k_rot = rope_at(k_vec, j) + dot_original = jnp.dot(q_rot, k_rot) + + q_rot_shifted = rope_at(q_vec, i + d) + k_rot_shifted = rope_at(k_vec, j + d) + dot_shifted = jnp.dot(q_rot_shifted, k_rot_shifted) + + np.testing.assert_allclose(dot_original, dot_shifted, atol=1e-5) + + +if __name__ == '__main__': + absltest.main() From 09a286068824c05d1b87ede0f40f55128c0caf61 Mon Sep 17 00:00:00 2001 From: Sam Anklesaria Date: Tue, 9 Jun 2026 16:10:05 -0600 Subject: [PATCH 2/2] Add multimodal llm example --- docs_nnx/examples/multimodal_llm.ipynb | 579 +++++++++++++++++++++++++ docs_nnx/examples/multimodal_llm.md | 375 ++++++++++++++++ 2 files changed, 954 insertions(+) create mode 100644 docs_nnx/examples/multimodal_llm.ipynb create mode 100644 docs_nnx/examples/multimodal_llm.md diff --git a/docs_nnx/examples/multimodal_llm.ipynb b/docs_nnx/examples/multimodal_llm.ipynb new file mode 100644 index 000000000..3f22c058e --- /dev/null +++ b/docs_nnx/examples/multimodal_llm.ipynb @@ -0,0 +1,579 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "61dd15b5", + "metadata": {}, + "source": [ + "# Build a multimodal language model with JAX\n", + "\n", + "This tutorial extends [miniGPT](minigpt.md) to handle images alongside text. You will build a small multimodal language model that learns to caption MNIST digits by projecting image patches directly into the token sequence — the same core idea used in [LLaVA](https://arxiv.org/abs/2304.08485).\n", + "\n", + "The model uses **Multimodal Rotary Position Embeddings (M-RoPE)** from [Qwen2-VL](https://arxiv.org/abs/2409.12191): image tokens are assigned independent 2D spatial positions (row, col) while text tokens get standard 1D sequential positions, with no separate learned positional embedding table.\n", + "\n", + "Here, you will learn how to:\n", + "\n", + "- Turn image patches into token embeddings with a linear patch projector\n", + "- Assign 2D spatial positions to image tokens and 1D positions to text tokens with M-RoPE\n", + "- Build a causal transformer that jointly attends over image and text tokens\n", + "- Train the model to generate image captions from scratch\n", + "\n", + "## Setup" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "1dba72ae", + "metadata": {}, + "outputs": [], + "source": [ + "import jax\n", + "import jax.numpy as jnp\n", + "import flax.nnx as nnx\n", + "import optax\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "from tensorboardX import SummaryWriter\n", + "import random\n", + "import functools\n", + "import tiktoken\n", + "from datasets import load_dataset" + ] + }, + { + "cell_type": "markdown", + "id": "12687e5b", + "metadata": {}, + "source": [ + "## Tokenizer\n", + "\n", + "We use the same GPT-2 tokenizer from [Tiktoken](https://github.com/openai/tiktoken) as in the miniGPT tutorial, giving the model a full-sized vocabulary consistent with a real language model.\n", + "\n", + "To make training more interesting, each image is paired with a caption drawn at random from several templates." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "e761d729-6f61-427b-8989-df81d143c70a", + "metadata": {}, + "outputs": [], + "source": [ + "tokenizer = tiktoken.get_encoding(\"gpt2\")\n", + "BOS = EOS = tokenizer.encode('<|endoftext|>', allowed_special={'<|endoftext|>'})[0]\n", + "PAD = 0\n", + "VOCAB_SIZE = tokenizer.n_vocab" + ] + }, + { + "cell_type": "markdown", + "id": "9dec8b87", + "metadata": {}, + "source": [ + "## Dataset: MNIST with captions\n", + "\n", + "We pair every digit image with a templated caption such as `\"the digit is seven\"`. The model learns to generate the caption given the image." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "f82599c3", + "metadata": {}, + "outputs": [], + "source": [ + "mnist = load_dataset(\"ylecun/mnist\")\n", + "x_train = np.stack([np.array(img) for img in mnist['train']['image']]).astype(np.float32) / 255.0\n", + "y_train = np.array(mnist['train']['label'])\n", + "x_test = np.stack([np.array(img) for img in mnist['test']['image']]).astype(np.float32) / 255.0\n", + "y_test = np.array(mnist['test']['label'])\n", + "\n", + "DIGIT_NAMES = ['zero', 'one', 'two', 'three', 'four',\n", + " 'five', 'six', 'seven', 'eight', 'nine']\n", + "\n", + "CAPTION_TEMPLATES = [\n", + " \"the digit is {word}\",\n", + " \"this shows the number {word}\",\n", + " \"I see a {word}\",\n", + " \"the handwritten digit {word}\",\n", + " \"a picture of {word}\",\n", + " \"this digit is {word}\",\n", + "]\n", + "\n", + "# TEXT_LEN must cover the longest possible encoded caption (BOS + tokens).\n", + "TEXT_LEN = max(\n", + " len([BOS] + tokenizer.encode(tmpl.format(word=d)))\n", + " for tmpl in CAPTION_TEMPLATES\n", + " for d in DIGIT_NAMES\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "3b8b2e7b", + "metadata": {}, + "outputs": [], + "source": [ + "def encode_caption(label: int, template: str | None = None) -> list[int]:\n", + " \"\"\"Encode a digit label using a randomly chosen (or specified) caption template.\"\"\"\n", + " tmpl = template or random.choice(CAPTION_TEMPLATES)\n", + " tokens = [BOS] + tokenizer.encode(tmpl.format(word=DIGIT_NAMES[label]))\n", + " return tokens + [PAD] * (TEXT_LEN - len(tokens))\n", + "\n", + "def decode_tokens(ids) -> str:\n", + " valid = [int(i) for i in ids if int(i) not in (PAD, BOS, EOS)]\n", + " return tokenizer.decode(valid)" + ] + }, + { + "cell_type": "markdown", + "id": "8a05f2a4", + "metadata": {}, + "source": [ + "The training objective is **next-token prediction** on the text part of the sequence.\n", + "Given the image and all previous text tokens, the model predicts the next token:\n", + "\n", + "| Input | `` | `the` | `digit` | `is` | `seven` |\n", + "|--------|---------|---------|---------|---------|---------|\n", + "| Target | `the` | `digit` | `is` | `seven` | `` |\n", + "\n", + "The loss is computed only over the text positions — not over the image patches. " + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "b37b2935", + "metadata": {}, + "outputs": [], + "source": [ + "def make_batch(images: np.ndarray, labels: np.ndarray):\n", + " \"\"\"Build (images, text_inputs, text_targets) JAX arrays for a batch.\"\"\"\n", + " B = len(labels)\n", + " text_in = np.array([encode_caption(lbl) for lbl in labels], dtype=np.int32)\n", + " text_out = np.concatenate(\n", + " [text_in[:, 1:], np.full((B, 1), EOS, dtype=np.int32)], axis=1\n", + " )\n", + " return jnp.array(images), jnp.array(text_in), jnp.array(text_out)" + ] + }, + { + "cell_type": "markdown", + "id": "899cd64e", + "metadata": {}, + "source": [ + "## Model architecture\n", + "\n", + "The model has four components assembled into `MultimodalGPT`:\n", + "\n", + "1. **`PatchEmbedding`** — splits each 28×28 image into 16 non-overlapping 7×7 patches and projects each flattened patch to `embed_dim` with a single linear layer.\n", + "2. **Token embedding** — a standard `nnx.Embed` lookup that maps text token IDs to `embed_dim` vectors.\n", + "3. **Causal transformer** — a stack of pre-norm transformer blocks. Each block uses `nnx.MultiHeadAttention` with a custom M-RoPE `attention_fn` so image patches and text tokens share the same sequence but carry different positional coordinates.\n", + "4. **LM head** — a linear projection from `embed_dim` to `vocab_size`, applied only to the text positions of the output. The image patch outputs are discarded because the loss is computed purely on text.\n", + "\n", + "The 16 image tokens are prepended to the text tokens and the combined sequence is processed jointly by the transformer. M-RoPE lets the model distinguish image patches by their 2D spatial location and text tokens by their sequential order.\n", + "\n", + "### Patch embedding\n", + "\n", + "We divide each 28×28 image into a 4×4 grid of 7×7 patches, yielding **16 image tokens**. A single linear layer projects each flattened patch to `embed_dim`." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "cac9bfa1", + "metadata": {}, + "outputs": [], + "source": [ + "class PatchEmbedding(nnx.Module):\n", + " \"\"\"Project image patches into the model's embedding space.\"\"\"\n", + "\n", + " def __init__(self, patch_size: int, embed_dim: int, *, rngs: nnx.Rngs):\n", + " self.patch_size = patch_size\n", + " self.proj = nnx.Linear(patch_size * patch_size, embed_dim, rngs=rngs)\n", + "\n", + " def __call__(self, images):\n", + " # images: (B, H, W)\n", + " B, H, W = images.shape\n", + " p = self.patch_size\n", + " nh, nw = H // p, W // p\n", + " x = images.reshape(B, nh, p, nw, p).transpose(0, 1, 3, 2, 4) # (B, nh, nw, p, p)\n", + " return self.proj(x.reshape(B, nh * nw, p * p)) # (B, num_patches, embed_dim)" + ] + }, + { + "cell_type": "markdown", + "id": "fd34ae63", + "metadata": {}, + "source": [ + "### Multimodal Rotary Position Embedding (M-RoPE)\n", + "\n", + "Standard RoPE assigns a single sequential integer to each token. M-RoPE extends this for multimodal sequences by **splitting the attention head dimension into two equal halves**:\n", + "\n", + "- **First half** is rotated by the *row* index (image patches) or *sequence position* (text).\n", + "- **Second half** is rotated by the *column* index (image patches) or *sequence position* (text).\n", + "\n", + "For text tokens both halves receive the same value, so M-RoPE reduces to standard 1D RoPE. For image patches the two halves carry independent row and column coordinates, encoding 2D location without any additional embedding table.\n", + "\n", + "```\n", + "Sequence layout: [patch(0,0) patch(0,1) … patch(3,3) | the digit is ]\n", + " ←————————————16 image tokens—————————→ ←————5 text tokens————→\n", + "\n", + "M-RoPE pos_0: [0 0 0 0 1 1 1 1 2 2 2 2 3 3 3 3 | 4 5 6 7 8] (row / seq)\n", + "M-RoPE pos_1: [0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 | 4 5 6 7 8] (col / seq)\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "85d15b38", + "metadata": {}, + "outputs": [], + "source": [ + "def make_mrope_positions(grid_h: int, grid_w: int, text_len: int):\n", + " \"\"\"Return per-token M-RoPE position IDs as two 1-D arrays.\n", + "\n", + " Image patches get independent (row, col) positions. Text tokens get a\n", + " sequential position in both components, offset past the image grid so\n", + " image and text coordinates do not collide.\n", + " \"\"\"\n", + " rows = jnp.repeat(jnp.arange(grid_h), grid_w)\n", + " cols = jnp.tile(jnp.arange(grid_w), grid_h)\n", + " text_seq = jnp.arange(text_len) + max(grid_h, grid_w)\n", + " return jnp.concatenate([rows, text_seq]), jnp.concatenate([cols, text_seq])" + ] + }, + { + "cell_type": "markdown", + "id": "a088f566", + "metadata": {}, + "source": [ + "### Attention, transformer blocks, and the full model\n", + "\n", + "M-RoPE is implemented with `dot_product_attention_with_mrope`, a thin extension of\n", + "`nnx.dot_product_attention_with_rope` for 2D positions. It accepts `input_positions`\n", + "as a `[2, seq_len]` array where row 0 holds row/sequence indices and row 1 holds\n", + "column/sequence indices. The first and second halves of each head's dimension are\n", + "rotated independently by `nnx.RoPE`, then concatenated before the standard\n", + "dot-product attention is computed. `nnx.MultiHeadAttention` takes this function as\n", + "its `attention_fn`, so all Q/K/V projections and the output projection are handled\n", + "by the built-in module." + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "id": "6a5c7589", + "metadata": {}, + "outputs": [], + "source": [ + "def dot_product_attention_with_mrope(\n", + " query, key, value, *, rope, input_positions, **kwargs\n", + "):\n", + " \"\"\"Dot-product attention with Multimodal RoPE.\n", + "\n", + " Extends nnx.dot_product_attention_with_rope for 2D spatial positions.\n", + " ``input_positions`` must be a [2, seq_len] array: row 0 contains row/sequence\n", + " indices (pos_0) and row 1 contains column/sequence indices (pos_1). Image\n", + " patches receive independent (row, col) coordinates; text tokens receive the\n", + " same value in both rows (standard 1-D RoPE behaviour).\n", + " \"\"\"\n", + " pos_0, pos_1 = input_positions[0], input_positions[1]\n", + " h = query.shape[-1] // 2\n", + " apply = jax.vmap(rope, in_axes=(-2, None), out_axes=-2)\n", + " query = jnp.concatenate([apply(query[..., :h], pos_0), apply(query[..., h:], pos_1)], axis=-1)\n", + " key = jnp.concatenate([apply(key [..., :h], pos_0), apply(key [..., h:], pos_1)], axis=-1)\n", + " return nnx.dot_product_attention(query, key, value, **kwargs)" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "f2edb6be", + "metadata": {}, + "outputs": [], + "source": [ + "class CausalSelfAttention(nnx.Module):\n", + " \"\"\"Causal multi-head self-attention with M-RoPE via nnx.MultiHeadAttention.\"\"\"\n", + "\n", + " def __init__(self, embed_dim: int, num_heads: int, *, rngs: nnx.Rngs):\n", + " assert embed_dim % num_heads == 0\n", + " assert (embed_dim // num_heads) % 4 == 0, \"head_dim must be divisible by 4 for M-RoPE\"\n", + " self.attn = nnx.MultiHeadAttention(\n", + " num_heads=num_heads,\n", + " in_features=embed_dim,\n", + " use_bias=False,\n", + " decode=False,\n", + " attention_fn=functools.partial(dot_product_attention_with_mrope, rope=nnx.RoPE()),\n", + " rngs=rngs,\n", + " )\n", + "\n", + " def __call__(self, x, pos_0, pos_1):\n", + " return self.attn(x, is_causal=True, input_positions=jnp.stack([pos_0, pos_1]))" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "401ae1de", + "metadata": {}, + "outputs": [], + "source": [ + "class TransformerBlock(nnx.Module):\n", + " \"\"\"Pre-norm transformer block.\"\"\"\n", + "\n", + " def __init__(self, embed_dim: int, num_heads: int, ff_dim: int, *, rngs: nnx.Rngs):\n", + " self.attn = CausalSelfAttention(embed_dim, num_heads, rngs=rngs)\n", + " self.fc1 = nnx.Linear(embed_dim, ff_dim, rngs=rngs)\n", + " self.fc2 = nnx.Linear(ff_dim, embed_dim, rngs=rngs)\n", + " self.norm1 = nnx.LayerNorm(num_features=embed_dim, rngs=rngs)\n", + " self.norm2 = nnx.LayerNorm(num_features=embed_dim, rngs=rngs)\n", + "\n", + " def __call__(self, x, pos_0, pos_1):\n", + " x = x + self.attn(self.norm1(x), pos_0, pos_1)\n", + " x = x + self.fc2(jax.nn.gelu(self.fc1(self.norm2(x))))\n", + " return x" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "90769405", + "metadata": {}, + "outputs": [], + "source": [ + "class MultimodalGPT(nnx.Module):\n", + " \"\"\"Image-conditioned causal language model.\n", + "\n", + " Image patches and text tokens are concatenated into a single sequence.\n", + " M-RoPE encodes spatial position for patches and sequential position for\n", + " text.\n", + " \"\"\"\n", + "\n", + " def __init__(self, *, rngs: nnx.Rngs):\n", + " self.patch_embed = PatchEmbedding(PATCH_SIZE, EMBED_DIM, rngs=rngs)\n", + " self.token_embed = nnx.Embed(num_embeddings=VOCAB_SIZE, features=EMBED_DIM, rngs=rngs)\n", + " self.blocks = nnx.List([TransformerBlock(EMBED_DIM, NUM_HEADS, FF_DIM, rngs=rngs)\n", + " for _ in range(NUM_BLOCKS)])\n", + " self.norm = nnx.LayerNorm(num_features=EMBED_DIM, rngs=rngs)\n", + " self.lm_head = nnx.Linear(EMBED_DIM, VOCAB_SIZE, use_bias=False, rngs=rngs)\n", + "\n", + " def __call__(self, images, text_tokens):\n", + " \"\"\"\n", + " Args:\n", + " images: (B, H, W) float32 in [0, 1]\n", + " text_tokens: (B, text_len) int32\n", + " Returns:\n", + " logits: (B, text_len, vocab_size)\n", + " \"\"\"\n", + " text_len = text_tokens.shape[1]\n", + " x = jnp.concatenate([\n", + " self.patch_embed(images), # (B, num_patches, embed_dim)\n", + " self.token_embed(text_tokens), # (B, text_len, embed_dim)\n", + " ], axis=1)\n", + "\n", + " pos_0, pos_1 = make_mrope_positions(GRID_H, GRID_W, text_len)\n", + " for block in self.blocks:\n", + " x = block(x, pos_0, pos_1)\n", + "\n", + " return self.lm_head(self.norm(x)[:, GRID_H * GRID_W:]) # text positions only\n", + "\n", + " def caption(self, image):\n", + " \"\"\"Greedily decode a caption for a single (H, W) image.\"\"\"\n", + " tokens = [BOS]\n", + " image = image[None] # add batch dim\n", + " for _ in range(TEXT_LEN - 1):\n", + " text_in = jnp.array([tokens + [PAD] * (TEXT_LEN - len(tokens))])\n", + " next_id = int(jnp.argmax(self(image, text_in)[0, len(tokens) - 1]))\n", + " if next_id == EOS:\n", + " break\n", + " tokens.append(next_id)\n", + " return decode_tokens(tokens[1:]) # strip BOS" + ] + }, + { + "cell_type": "markdown", + "id": "30d31950", + "metadata": {}, + "source": [ + "## Set hyperparameters and create the model" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "f01780e1", + "metadata": {}, + "outputs": [], + "source": [ + "PATCH_SIZE = 7 # 28 / 7 = 4 patches per side\n", + "GRID_H = GRID_W = 28 // PATCH_SIZE\n", + "EMBED_DIM = 64\n", + "NUM_HEADS = 4 # head_dim = 16, divisible by 4 ✓\n", + "FF_DIM = 256\n", + "NUM_BLOCKS = 4\n", + "BATCH_SIZE = 256\n", + "NUM_EPOCHS = 5\n", + "LR = 1e-3" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "id": "fd8b5b25", + "metadata": {}, + "outputs": [], + "source": [ + "model = MultimodalGPT(rngs=nnx.Rngs(0))\n", + "optimizer = nnx.Optimizer(model, optax.adam(LR), wrt=nnx.Param)" + ] + }, + { + "cell_type": "markdown", + "id": "86dd2dc5", + "metadata": {}, + "source": [ + "## Define the loss function and training step\n", + "\n", + "The training objective is **next-token prediction** on the text part of the sequence. `loss_fn` runs a forward pass and computes the mean cross-entropy between the model's logits (shape `[B, text_len, vocab_size]`) and the shifted target token IDs from `make_batch`. Image patch positions are excluded from the loss — the model only needs to learn to predict text given the visual context.\n", + "\n", + "`train_step` wraps `loss_fn` with `nnx.value_and_grad` to obtain both the scalar loss and the parameter gradients in one call, then applies the update through the Optax optimizer. Decorating the step with `@nnx.jit` compiles the entire forward-backward-update pass with XLA so it runs efficiently on GPU and TPU. " + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "id": "09e2f468", + "metadata": {}, + "outputs": [], + "source": [ + "def loss_fn(model, images, text_in, text_out):\n", + " logits = model(images, text_in)\n", + " return optax.softmax_cross_entropy_with_integer_labels(logits, text_out).mean()\n", + "\n", + "@nnx.jit\n", + "def train_step(model, optimizer, images, text_in, text_out):\n", + " loss, grads = nnx.value_and_grad(loss_fn)(model, images, text_in, text_out)\n", + " optimizer.update(model, grads)\n", + " return loss" + ] + }, + { + "cell_type": "markdown", + "id": "c75c174c", + "metadata": {}, + "source": [ + "## Train the model" + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "id": "501b3b55", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Epoch 1/5 loss=1.6708\n", + "Epoch 2/5 loss=0.3624\n", + "Epoch 3/5 loss=0.3343\n", + "Epoch 4/5 loss=0.3239\n", + "Epoch 5/5 loss=0.3189\n" + ] + } + ], + "source": [ + "writer = SummaryWriter()\n", + "step = 0\n", + "for epoch in range(NUM_EPOCHS):\n", + " perm = np.random.permutation(len(x_train))\n", + " epoch_losses = []\n", + "\n", + " for i in range(0, len(x_train) - BATCH_SIZE + 1, BATCH_SIZE):\n", + " idx = perm[i : i + BATCH_SIZE]\n", + " images, text_in, text_out = make_batch(x_train[idx], y_train[idx])\n", + " loss = float(train_step(model, optimizer, images, text_in, text_out))\n", + " writer.add_scalar('train_loss', loss, step)\n", + " epoch_losses.append(loss)\n", + " step += 1\n", + "\n", + " print(f\"Epoch {epoch + 1}/{NUM_EPOCHS} loss={np.mean(epoch_losses):.4f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "0645116b", + "metadata": {}, + "source": [ + "## Evaluate" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "id": "237ea51c", + "metadata": {}, + "outputs": [ + { + "data": { + "image/png": "iVBORw0KGgoAAAANSUhEUgAABK4AAAHxCAYAAABEevT6AAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAPYQAAD2EBqD+naQAAbrNJREFUeJzt3XmcjXX/x/H3mcUyg5kxZpBlbC22yB4GlVJESKWVEkpa7/aFUShUlEK675AoUpYKSSGyhUz2IkZk11jGzGDm+v1xfnM4Zs53Zs45M3MNr+fj0eO+Xe/z/V7f65zre10zn7mu6zgsy7IEAAAAAAAA2ExAQQ8AAAAAAAAAyAqFKwAAAAAAANgShSsAAAAAAADYEoUrAAAAAAAA2BKFKwAAAAAAANgShSsAAAAAAADYEoUrAAAAAAAA2BKFKwAAAAAAANgShSsAAAAAAADYkt8KV45BDvWf299f3RX4enKjyqgq6jmrZ56uo83ENmozsU2eriOvLd61WI5BDi3etdi1rOesnqoyqopX/cUtjpNjkMM/g7uE5XROTVw/UY5BDu1K3JXnY9qVuEuOQQ69vfztPF9Xfst4H9f8s6agh5LvLubPNT8xZ/MXc/bi/FzzGvM0fzFPL87PtSAwd/MXc7fgP9cDJw+o2/RuihweKccgh0atHFWg4/EkV4Wr5X8vV9ziOCWmJObRcJBT/5z4R3GL47R+//pM2dQNU227w+W1oUuHatbWWQU9DNth7ua/Mb+O0cT1Ewt6GLiA6dhpJ8zZ/MecRW4xT/Mf8xT+wNzNf8xd+3r6+6f1/Y7v9VLLlzS5y2TdXOPmgh5SlnJduBq0ZBCTvAAsuH+BFty/wPXvf078o0FLBhX6wtXHHT/Wtv7bvGr7aqtXlfxKstsyCldZ88fcvf/q+5X8SrJiwmL8N7CLGCdoezIdO+2EOZv/mLPILeZp/mOewh+Yu/mPuWtfP+38SbddeZuebf6s7rv6Pl1V5qqCHlKWeMaVzZ06c0qSVCSwiIoEFing0fhfcGCwigYV9aptUECQigUV8/OI4ElgQKCKBRWTw8HtmbC3pNNJBT0EW2DOorC4lOcs8xSFRUHO05SzKUq30gts/Vlh7qKwyG7uHkw6qPBi4fkzmAtk1DpyIseFq7jFcXruh+ckSVXfqyrHIEeW9/XO2jpLdcbUUdHBRVV7TG3N3z4/U197j+/VQ7MfUtm3y7pe98lvn+R40DlZT0Jigvp9109XfnClig8prsjhkbrjyzsyjTfjvtpfdv+iZ75/RlEjohQ6NFRdpnXRoaRDbq+1LEuDfx6siu9WVMiQEF036TptOrjJ7TWJKYkKfD1Q769637Xs8KnDChgUoMjhkbIsy7X80W8fVbm3y7n+3WZiG9UZU0dr/1mrVhNaKWRIiF7+8WVXlvGMq8W7Fqvxx40lSQ/OftD1WUxcP1FtJrbRd39+p4RjCa7l5z9DKvVsqgYuGqga79dQ0cFFVWlkJT3/w/NKPZvqth0Z93fn5PPMyp7je9T5i84KHRqq6BHRenr+05nWIWX9jKsjp47o/pn3q9SbpRT+Vrh6zOqh+P3xrm3McOEzrhyDHEo6k6RJ8ZNc257ds8dGrxqt2mNqK2RIiCKGRajR+EaaumGq22uy218PnDygoNeDNGjxoEz9bzu8TY5BDn2w+gPXssSURD01/ylVGllJRQcXVY33a2jYsmFuJ+Tz73kev3a8qr9fXUUHF1Xjjxvr172/GrfpQv6au1ndy7/mnzVq91k7lRleRsWHFFfV96rqodkPZTum3LTLyfb/tPMnxU6IVejQUIW/Fa7bvrhNWw5tceW/H/hdjkEOzdk2x7Vs7T9r5RjkUIOPGrj1dcuUW9T0v0192sYqo6po06FNWpKwxPV+X/iMutSzqdkecyRp3p/zXNtW8s2S6jC1Q6bjTlZyc2xzDHIobnFclttx/hzK6HPZ7mV6Yt4TihoRpfC3wtX3m746nXZaiSmJemDmA4oYFqGIYRF6/ofn3Y555xu5YqRiRsWo+JDiaj2xtTYe3JjpNVsPb1W36d1UelhpFRtcTI3GN3L7DM8f05JdS9Tvu36KHhGtiiMrZrlO07Hz/VXvK/D1QLe/vr6z/B05Bjn0zPfPuJalpaep5Jsl9cIPL7iWJZ1O0n++/49rTl/5wZV6e/nbHrc9O8xZ5ixz1t3BpIPqNbuXyr5dVsUGF1O9cfU0af0kt9fk9ryZk7GaME+Zp8xT923K+Ewu/O/8Z9vm5HfAjGfifrHxC73606uq8G4FhQwJ0fHU45KkLzd9qYbjG6r4kOIqM7yM7vv6Pu09vtfj2C7E3GXuMnfdX2vJ0oe/fuj6bDL89e9fuuPLO1R6WGmFDAlRs/8203d/fJdlHxfOn6yebW2qdeREUE5f2LVmV/1x5A99vvFzjWw3UmVCykiSokKiXK9ZtnuZvt7ytfo17qeSRUrq/dXv6/bpt2v3U7sVGRIpyflLfrP/NZNDDvVv3F9RoVGat32ees3ppeOpx/VUs6eyHUtO1vPrP79q+d/L1b12d1UsVVG7Endp7JqxajOxjTY/tlkhwSFufT4+73FFFI/QwNYDtStxl0atHKX+gf01rds012sGLBqgwUsHq/3l7dW+Rnut27dON312k06nnXa9JrxYuOpE19HPCT/riaZPuMbrcDh0NPmoNh/arNrRtSVJS3cvVWxMrNs4jiQf0S1TblH3Ot1139X3qWxo2UzbX7NMTb3e5nUNWDxAfRr0cfXRvFJzVShZQcdSj2nP8T0a2W6kJKlEkRKSpHQrXZ2+6KRlu5epT4M+qhlVUxsObNDIlSP1x5E/NKv7rFy/z1lJPpOsGz69QbuP7dYTTZ7QZSUv0+TfJ+unXT95/lD/X7qVro6fd9Tqvav1aKNHdVWZqzR722z1mNUj27aTu0zWw3MeVpMKTdSnYR9JUvWI6h5f//Haj/XE/CfUrVY3Pdn0SaWcTdHvB37Xqj2rdE/deyTlbH8tW6KsWldprembp2tgm4Fu65i2aZoCHYG6o9YdkpxV5dYTW2vv8b3q27CvKodV1vI9y/XSjy9p38l9GnXzKLf2UzdM1YnTJ9S3YV855NDw5cPVdXpX/fXEXwoODM72PZH8N3cvdDDpoG6afJOiQqP0YssXFV4sXLsSd+nrLV8bx5ObdjnZ/oV/LdQtU25RtYhqimsdp+SzyRq9erRafNJC6/quU5XwKqoTXUfhxcL1c8LP6nRlJ0nO+RfgCFD8gXgdTz2uUkVLKd1K1/K/l6tPgz4+beOom0fp8XmPq0SREnol9hVJyjSXc3LMmRw/WT1m9VC7Gu00rO0wnTpzSmPXjFXLCS31W9/fVCW8inEcOV1Pbj0+73GVK1FOg9oM0so9KzV+3XiFFwvX8j3LVTmssoZeP1Rzt8/ViOUjVCe6jh6o94Bb+0/jP9WJ0yf0WOPHlHI2Re+tek/XT7peGx7doLIlnO/TpoOb1OKTFqpQqoJebPmiQoNDNX3zdHX+orO+uvMrdanZxa3PfnP7KSokSgNaD/D4lyXTsfNYyjGlW+latnuZbr3iVknn9pGlu5e6+vht/286efqkWsW0kuT8g0anLzpp0c5F6nVNL9UvV1/f7/hez/3wnPYe36uRN4/M9fvLnGXOMmfPST6TrDYT22j70e3q36S/qoZX1Zebv1TP2T2VmJKoJ5s96fb6nOyDuR1rVpinzFPm6Tmjbh6lk6dPui0buXKk1u9fr8ji3v0O+MbPb6hIYBE9e+2zSk1LVZHAIpq4fqIenP2gGl/WWG/e8KYOnDyg91a9p1/+/kW/9f0tR1eNMHeZu8xdp1YxrTS5y2TdP/N+3VjtRrdxHTh5QM3/11ynzpzSE02fUGTxSE2Kn6ROX3TSjDtm5Og8mZWc1Do8snJhxC8jLMXJ2vnvzkyZ4mQVeaOItf3Idtey+P3xluJkjV412rWs1+xeVvm3y1uHkw67te8+o7sV9maYder0KeMYcrqerPpZ8fcKS3GyPl3/qWvZhN8mWIqT1fbTtlZ6erpr+dPzn7YCBwVaicmJlmVZ1sGTB60ibxSxOkzp4Pa6lxe+bClOVo+ZPVzLHvvuMavsiLKufz8z/xmr1YRWVvSIaGvsr2Mty7KsI6eOWI44h/Xeyvdcr2s9obWlOFnjfh2XaeytJ7S2Wk9o7fr3r3t/tRQna8JvEzK9tsOUDlbMyJhMyyfHT7YCBgVYSxOWui0f9+s4S3Gyftn9i2tZTt/nrIxaMcpSnKzpG6e7liWdTrJqvF/DUpysRTsXuZb3mNnDbaxfbf7KUpysUStGuZalpadZ10+6PtP2Dlw00FKc+y4cOiTU7bMwue3z26zaH9Y2vian++tHaz6yFCdrw4ENbq+r9WEt6/pJ17v+/caSN6zQIaHWH4f/cHvdiz+8aAUOCrR2J+62LMuydv6701KcrMhhkdbRU0ddr5u9dbalOFnfbPsmR9uYwR9zN2OuZPQxc8tMS3Gyft37a67GkpN2udn++uPqW9Ejoq0jp464jT9gUID1wMwHXMs6TOlgNfm4ievfXad1tbpO62oFDgq05v05z7Isy1r3zzpLcbJmb53t0zZalmXV/rC225zNkNNjzonUE1b4W+FW7zm93drvP7HfCnszLNNyb9djWc59YOCigZn6iBkZ4zafMvpsN7mdW5/X/vdayxHnsB755hHXsrNpZ62K71Z0ew8yPtfig4tbe47tcS1ftWeVpThZT89/2rXshkk3WHXH1LVSzqS4lqWnp1vN/9fcuvz9yzONqeUnLa2zaWeN74lleT52pqWnWaXeLGU9v+B517oih0Vad0y/wwocFGidSD1hWZZlvbv8XStgUID1b/K/lmVZ1qwtsyzFyRq8ZLBbf92md7MccQ63eZUbzFnmrKf1WNalNWczfqb4LP4z17LTZ09b1/73WqvE0BLW8ZTjbmPNyT6Y07Fmh3nKPPW0Hsu6tObphaZvnG4pTtbri193Lcvpz9SLdi6yFCer2nvV3H6fO332tBU9ItqqM6aOlXwm2bX8223fWoqTNeCnATkeH3OXuetpPZZ16c1dxcl67LvH3JY9Ne8pS3FyqxucSD1hVR1V1aoyqoqVlp7mtr4L51LGPD7/935TrSMnArwqlXnQtlpbVS997gqXq8terVJFS+mvf//KKJLpqy1fqeMVHWXJ0uFTh13/taveTsdSj2ndvnU+r0eSigcXd/3/M2lndOTUEdUoXUPhxcKzXEefBn3c7lGOrRyrNCtNCccSJDmr16fTTuvxJo+7vS6rK8RiK8fqQNIBbTvsfOj40t1L1apyK8VWjnX95X7Z7mWyZCm2svsVV0UDi+rBax7M9j3wxpebv1TNMjV1VZmr3N7766teL0latHOR2+tz8j5nZe72uSpfory61ermWhYSHOKq1pvM3z5fwQHB6t2wt2tZgCNAjzV+LEfbmBvhxcK15/gej7fe5WZ/7Vqzq4ICgjRt47lq/caDG7X50GbdVfsu17IvN3+p2JhYRRSPcOuvbbW2SrPS9HPCz25juKv2XYooHuH6d8b+kt1nkFvefNYZf9X69o9vdSbtTI7XlZt22W3/vhP7tH7/evWs11Oli5d2G/+N1W7U3D/nurVdt2+d668Oy3YvU/sa7VW/XH0tTXDOy6W7l8ohh1pWbunTNuZEdsecH3b8oMSURN1d5263fSUwIFBNKzbVol2LPHWdq/V4o9c1vdz6bFqhqSxZ6tWgl2tZYECgGl3WKMt9qPNVnVWhVAXXv5tUaKKmFZq6Pq+jyUf1086fdGftO3Xi9AnXth9JPqJ21dvpz6N/ZrotoHeD3goMCPR6mwIcAWpeqbl+3u2cg1sOb9GR5CN6seWLsmRpxd8rJDn3kYy/VErS3D/nKtAR6LrCNsN/rv2PLFmat32e12MyYc4yZ3OjMM/ZudvnqlyJcrq77t2uZcGBwXqi6RM6efqkliQscXt9dvugN2P1FvOUeZobhXmenm/zoc16aM5Duu3K2/Rqq1clefc7YI96Pdx+n1vzzxodTDqofo36uT3jtsMVHXRVmav03Z/utzD5grnL3M2Ni2Xunm/u9rlqUqGJ67OXnHdx9WnYR7sSd2nzoc1e9etLrSPHtwrmROVSlTMtiygWoX9T/pUkHTp1SIkpiRq/brzGrxufZR8Hkw76vB7JeWn5m8ve1IT1E7T3+F5ZOnc/6bHUY5n7DHPvM+Og8G+ys8+MnfnyyMvdXhcVGqWIYhFuyzJuP1m6e6kqlqqo3/b/psHXD1ZUaJTeXv62M0tYqlJFS6leuXpubSuUqpBnD2H/88if2nJ4i6JGRGWZX/je5+R9zkpCYoJqlK6R6WGFV5a5MtsxJhxLUPmS5TPdylmjdI1s2+bWCy1e0MK/FqrJf5uoRukauqnaTbqn7j1qUbmFpNztr2VCyuiGqjdo+ubpeuP6NyRJ0zZOU1BAkLrW7Op6/Z9H/tTvB37P+WeQzX7pL9581q1jWuv2mrdr0JJBGrlypNpUaaPOV3bWPXXvMT5wPzftcjovs9q3apapqe93fK+k00kKLRKq2JhYnU0/qxV7VqhSqUo6mHRQsTGx2nRok6ugvHT3UtWKquU62Xu7jTmR3bb9efRPSdL1n16fZftSRUv5ZT3euLDPsGJhkqRKpSq5Ly8aluU+dHnpyzMtuyLyCk3fNF2StP3odlmy9Nqi1/TaoteyHMPBpINuJ/mq4VVztxFZiK0cq7jFcUo+k6ylCUtVvkR5NSjfQPXK1tPS3Ut1Y/UbtWz3Mt1Z+05Xm4RjCbqs5GUqWbSkW181y9R05one/yBkwpxlzuZGYZ6zCYkJurz05QpwuP+t1dMcy+7982as3mKeMk9zozDP0wzHU4+r67SuqlCygj7t8qnrdwFvfge8cN2m/fKqMldp2e5luRqrCXOXuZsbF8PcvVBCYoKa1m6aafn559460XVy3a8vtQ6/Fq48VfWs/38IWcbDp++7+j71qJf1M4uuLnu1z+uRnPeaTlg/QU81fUrXVrpWYUXD5HA41H1G9yy/lcJjn+cVvHLqspKXqWp4Vf2c8LOqhFeRZVm6tuK1igqJ0pPzn1RCYoKW7l6q5pWaZ/pBrHhQcQ+9+i7dSlfd6Lp6t927WeYXTq6cvM+FWc2omtrWf5u+/eNbzd8xX19t+Upj1ozRgFYDNOi6QbneX7vX6a4HZz+o9fvXq365+pq+ebpuqHqD6955yfkZ3FjtRj3f4vks+7si8gq3f/tzvzTx5rN2OByacecMrdyzUt9s+0bf7/heD815SO+seEcrH17peraaL+38uf2NLmukYkHF9HPCz6ocVlnRodG6IvIKxVaO1Zhfxyj1bKqWJixVl6vO3bPt7TbmRHbblrH/Te4yWeVKlMv0uqCAnB2+fXkP06y0XPWZ1XJvjhcZ2/7stc+qXY12Wb7mwmL2+X+V9VbLyi11Jv2MVuxZ4fYMwoyrZbce3qpDpw5lulK2IDBnmbNZudTmbFZy+jnlZqx+HwvzlHmaiz4L0zztOaun/jnxj1b3Xu1WSPDmd8C8OkbkBHOXuZuVi3nuesuhrL9VMy096/fKl1pHrgpXngaWU1EhUSpZpKTS0tPUtlpbn/rKzozNM9SjXg+90+4d17KUsylu3xiVGzFhMZKcV8xUi6jmWn4o6VCWldPYmFj9nPCzqoZXVf1y9VWyaEnVK1dPYUXDNH/7fK3bt06D2mT+FrqcMn0Wnr6WtXrp6orfH68bqt6Qp1/dGhMeo40HN8qyLLf1ZNw6aWwbFqNFOxfp1JlTblddbT+6PUfrzu12hRYJ1V117tJdde7S6bTT6jqtq4YsHaKXYl/K9f7a+arO6vttX9ftgn8c+UMvtXzJ7TXVS1fXydMn83z/v5Cvc9ekWcVmalaxmYbcMERTN0zVvV/fqy82fqGHGzycJ+3OlzEvs9q3th7ZqjIhZRRaJFSSVCSwiJpUaKKlu5eqclhlV/EhNiZWqWmpmrJhig4kHXA9dNvXsfo6xzK+WCA6NDrP95eIYhGZjo2n005r34l9ebK+jL+ane+PI3+4HqqZcYwNDgz2+7ab5kKTCk1UJLCIliYs1dLdS/Vcc+e3DrWKaaWP132sH//60fXvDDFhMVr410KdSD3hdtXV1sNbnXl4jN/H6SvmbNaYs54V5JyNCY/R7wd+V7qV7vbHPm/nmD/HyjxlnppcSvNUkt5a9pZmbZ2lr+/6WleVucot88fvgOfvlxmPOcmw7fA2V54TzF3mrsmlNnezEhMeo21HstiPLjj3ZlyxduH75cutl57k6hlXGTu6t8WfwIBA3V7rdn215assv+Ixq6+89FZgQGCmyunoVaM9Vkqz07ZaWwUHBGv06tFuldJRK0dl+frYyrHalbhL0zZNcx0IMp6h8u7Kd3Um/UymbxTMDdNnERocmuXtkHfWulN7T+zVx+s+zpQln0k2fltIbrSv0V7/nPhHMzbPcC07deaUx0uDz9euejudST+jj9eeG2O6la4Pf/0wR+sODQ7N8f555NQRt38XCSyiWlG1ZMnSmbQzud5fw4uFq131dpq+ebq+2PiFigQWUeerOru95s5ad2rFnhX6fvv3mfpLTEnU2fSzORp7bvk6d7Pyb/K/mf5qUL9cfUnOr7T1d7uslC9ZXvXL1dek+Elu27bx4EYt2LFA7S9v7/b62MqxWrVnlRbtXOSal2VCyqhmmZoa9ssw52vOm5e+jDU3+2JW2tVop1JFS2no0qFZPkfAn8fL6qWrZ3q+2vi1470+XmZn1tZZbvfir967Wqv2rtItNW6R5PyhpE2VNvpo7UdZ/pDgy7ab5kKxoGJqfFljfb7xc+0+ttvth7jks8l6f/X7qh5RXeVLlne1aX95e6VZafpg9QdufY1cOVIOOVzb5M9xeos5y5z1VkHO2fY12mv/yf1uz5A8m35Wo1ePVokiJdQ6pnWu+vPnWJmnTszTrF1K83ThXwv16k+v6pXYVzL97Cv553fARpc1UnRotMatHee2n8z7c562HN6iDpd3yPF4mbtOzN2sXUpz15P2Ndpr9d7Vrme8SlLS6SSNXzdeVcKrqFZULUnnCornv19p6Wkavzb73/tzK1dXXDUs31CS9MpPr6h77e4KDgxWxys6uiZ/Trx1w1tatHORmv63qXo36K1aUbV0NPmo1u1bp4V/LdTRF47mbgs8uPWKWzU5frLCioapVlQtrdizQgv/Wuj6StbcigqN0rPNn9Wby97UrZ/fqvY12uu3/b9p3vZ5breCZciY/NuObNPQG4a6lreKaaV52+epaGBRNb6ssXcbJ+dOEl4sXOPWjFPJIiUVWiRUTSs0VdWIqmpYvqGmbZqmZ75/Ro0va6wSRUqo45UddX+9+zV983Q98u0jWrRrkVpUaqG09DRtPbxV0zdP1/f3fa9GlzXyekwZejfsrQ9+/UAPzHpAa/etVfkS5TX598mZnluVlc5XdVaTCk30nwX/0faj23VVmas05485Oprs3C+y+wtJw8saauFfC/Xuinddt2w2rZj5/lxJuumzm1SuRDm1qNRCZUPLasvhLfpg9QfqcHkH15UTud1f76p9l+6beZ/G/DpG7aq3y/S1vM+1eE5z/pijWz+/VT3r9VTDyxoq6XSSNhzcoBmbZ2jXU7uy3J985Y+5e6FJ8ZM05tcx6nJVF1UvXV0nUk/o43Ufq1TRUplOjP5o58mIG0folim36Nr/Xate1/RS8hnn1/6GFQ1TXOs4t9fGVo7VkKVD9Pfxv91OxK1iWumjtR+pSngVVSxV0S9jbVi+ocauGavBPw9WjdI1FB0anekvhCalipbS2A5jdf/M+9VgfAN1r91dUaFR2n1st7778zu1qNRCH7T/IPuOcuDhax7WI989otun364bq92o+P3x+n7H93myL0rOy5lbTmipRxs9qtSzqRq1apQii0e63UL7YfsP1fKTlqo7tq56N+itahHVdCDpgFbsWaE9x/co/pF4r9ZtOnZKzn3krV/eUljRMNUtW1eS8weGKyOv1LYj29Szfk+3/jpe2VHXVblOr/z0inYl7lK9cvW0YMcCzd42W081fcrt4a65wZx1Ys5m7VKas30a9tFHaz9Sz9k9tXbfWlUJr6IZm2fol79/0ah2ozI9Xy4n/DVW5qkT8zRrl9I8vfuruxUVGqXLIy/XZ79/5pbdWO1GlS1R1uffAYMDgzWs7TA9OPtBtZ7YWnfXuVsHkg7ovVXvqUp4FT197dM5Hi9z14m5m7VLae568mLLF/X5xs91y5Rb9ETTJ1S6eGlNip+knf/u1Fd3fuW6Arp2dG01q9hML/34ko4mH1Xp4qX1xcYv8uRijFwVrhpXaKw3rntD49aM0/zt85VupWvnkztzNcnLliir1b1X6/Ulr+vrLV9rzK9jFBkSqdpRtTWs7bBcb4An7938ngIdgZqyYYpSzqaoRaUWWnj/QrX7LOv7QnNi8PWDVSyomMatGec88FZsqgX3LVCHqZkr/FeWuVLRodE6mHTQ7Wn8GQWtJhWa+PQQu+DAYE3qPEkv/fiSHvnuEZ1NP6sJt01Q1Yiq6te4n9YfWK8J6ydo5MqRigmLUccrOyrAEaBZd83SyJUj9Wn8p5q5ZaZCgkNULaKanmz6ZKbnK3krJDhEPz7wox6f97hGrx6tkOAQ3Vv3Xt1S4xbdPOVmY9vAgEB9d893enL+k5oUP0kBjgB1qdlFA1sPVItPWrh9i0hW3r3pXfX5to9e/elVJZ9NVo96PTwWrvo27KspG6bo3RXv6uTpk6pYqqKeaPqE6xtQpNzvr52u7KTiQcV14vQJt28TPP+9WdJziYYuHaovN3+pT3//VKWKltIVkVdoUJtBCisaZtw+b/lj7l6odUxrrd67Wl9s+kIHTh5QWLEwNanQRFO6TnEVAfzZzpO21dpq/r3zNXDxQA1YNEDBgcFqHdNaw9oOy9Rf80rNFegIVEhwiOqVPffFCLGVY/XR2o8yPbvIl7EOaD1ACccSNPyX4Tpx+oRax7TO1Qlaku6pe48uK3mZ3lr2lkYsH6HUtFRVKFlBsTGxerC+/759tHfD3tqZuFP/++1/mr99vmIrx+qH+3/QDZ/e4Ld1nO+Beg8owBGgUStH6WDSQTWp0EQftP/A7UqmWlG1tKbPGg1aMkgT4yfqyKkjig6N1jXlr9GAVgO8Xrfp2Ck5/7L41i9vZXoGYWzlWG07si3TPhLgCNCcu+dowKIBmrZpmiasn6Aq4VU04sYR+s+1//F6nMxZJ+Zs1i6lOVs8uLgW91ysFxe+qEnxk3Q89biujLxSE26bkKmQnFP+Givz1Il5mrVLaZ4ePnVYktRjVubnVy3qsUhlS5T1y++APev3VEhwiN5a9pZeWPiCQouEqstVXTSs7bBMfyw2Ye46MXezdinNXU/Kliir5b2W64WFL2j06tFKOZuiq8terW/u/kYdrnCvfUzpOkV9v+2rt5a9pfBi4ep1TS9dV/U63Tj5Rr+OyWFdLE/axkVt1tZZ6jKti5Y9uMz1rX8AAAAAAODilqtnXAH5IflMstu/09LTNHr1aJUqWkoNyjcooFEBAAAAAID8lqtbBYH88Pi8x5V8NlnXVrxWqWdT9fXWr7X87+Uaev3QAv1qXAAAAAAAkL+4VRC2M3XDVL2z4h1tP7pdKWdTVKN0DT3a6FH1b9K/oIcGAAAAAADyEYUrAAAAAAAA2BLPuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC355VsFHQ6HP7oB4IG/H0XHnAXyFnMWKFyYs0DhwpwFChdf5yxXXAEAAAAAAMCWKFwBAAAAAADAlihcAQAAAAAAwJYoXAEAAAAAAMCWKFwBAAAAAADAlihcAQAAAAAAwJYoXAEAAAAAAMCWKFwBAAAAAADAlihcAQAAAAAAwJYoXAEAAAAAAMCWKFwBAAAAAADAlihcAQAAAAAAwJYoXAEAAAAAAMCWKFwBAAAAAADAlihcAQAAAAAAwJYoXAEAAAAAAMCWKFwBAAAAAADAlihcAQAAAAAAwJYoXAEAAAAAAMCWKFwBAAAAAADAlihcAQAAAAAAwJaCCnoAAFAYPPvss8a8ePHixvzqq6825t26dcv1mDKMHTvWmK9YscKYT5482et1AwAAAEBe4oorAAAAAAAA2BKFKwAAAAAAANgShSsAAAAAAADYEoUrAAAAAAAA2BKFKwAAAAAAANgShSsAAAAAAADYEoUrAAAAAAAA2JLDsizL504cDn+MBYAHfpimbpizWZs2bZrHrFu3bvk4Ev/asWOHMW/btq3HbPfu3f4eziWBOQtfXHHFFcZ869atHrMnn3zS2Hb06NFejelix5wt/EJDQ435iBEjPGZ9+/Y1tl27dq0xv+OOO4x5QkKCMUfuMWeBwsXXOcsVVwAAAAAAALAlClcAAAAAAACwJQpXAAAAAAAAsCUKVwAAAAAAALAlClcAAAAAAACwJQpXAAAAAAAAsCUKVwAAAAAAALCloIIeAADkl2nTphnzbt265dm6t27dasy///57j1m1atWMbTt27GjMq1evbszvvfdej9mbb75pbAvA/6655hpjnp6e7jHbs2ePv4cDFArly5c35r179/aYmeaUJDVs2NCY33rrrcb8ww8/NOZAYdSgQQNj/vXXXxvzKlWq+HE09nHTTTcZ8y1bthjzv//+25/DuWhwxRUAAAAAAABsicIVAAAAAAAAbInCFQAAAAAAAGyJwhUAAAAAAABsicIVAAAAAAAAbInCFQAAAAAAAGwpqKAHAAD+0qhRI2PepUsXr/vetGmTMe/UqZMxP3z4sDE/efKkx6xIkSLGtitXrjTm9erVM+aRkZHGHED+ql+/vjFPSkrymM2cOdPPowHsISoqyphPmjQpn0YCQJLatWtnzIsWLZpPI7GXjh07GvOHHnrImHfv3t2fw7locMUVAAAAAAAAbInCFQAAAAAAAGyJwhUAAAAAAABsicIVAAAAAAAAbInCFQAAAAAAAGyJwhUAAAAAAABsicIVAAAAAAAAbCmooAdQkLp162bMe/fubcz/+ecfY56SkuIxmzJlirHt/v37jfn27duNOXApKl++vDF3OBzGfNOmTR6zdu3aGdvu27fPmPviP//5jzGvVauWT/1/9913PrUHkDt16tQx5v379zfmkydP9udwAFt44oknjHnnzp2NeZMmTfw4mtxp1aqVMQ8I8HytQHx8vLHtzz//7NWYAH8ICvJcLmjfvn0+jqTwWLt2rTF/5plnjHloaKjHLCkpyasxXQy44goAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2FFTQAyhIw4cPN+ZVqlTJs3X37dvXmJ84ccKYb9q0yZ/DKTT27NljzE2f6Zo1a/w9HNjMN998Y8xr1KhhzE3z7ujRo16NyR+6d+9uzIODg/NpJAD84aqrrjLmoaGhxnzatGn+HA5gCyNHjjTm6enp+TSS3OvatavXeUJCgrHtXXfdZczXrl1rzAFfXHfddR6za6+91tg2u9+1L1YRERHGvFatWsY8JCTEY5aUlOTVmC4GXHEFAAAAAAAAW6JwBQAAAAAAAFuicAUAAAAAAABbonAFAAAAAAAAW6JwBQAAAAAAAFuicAUAAAAAAABbonAFAAAAAAAAWwoq6AEUpN69exvzq6++2phv2bLFmNesWdNj1qBBA2PbNm3aGPNmzZoZ87///ttjVqlSJWNbX509e9ZjdujQIWPb8uXL+7Tu3bt3e8zWrFnjU98o/BISEgp6CB4999xzHrMrrrjCp75XrVrlUw7Av55//nljnt2xivMZCqu5c+d6zAIC7Pv39CNHjhjzkydPGvOYmBiPWdWqVY1tV69ebcwDAwONOWBSp04dY/755597zHbs2GFsO3ToUK/GVNjddtttBT2Ei5J9zxAAAAAAAAC4pFG4AgAAAAAAgC1RuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC1RuAIAAAAAAIAtBRX0AArSjz/+6FOenfnz53vdNiIiwpjXr1/fmK9du9Zj1rhxY2+GlGMpKSkesz/++MPYdsuWLca8dOnSxjy7r2UFCsqtt95qzF9//XWPWZEiRYxtDx48aMxfeuklY37q1CljDiB3qlSpYswbNWpkzLM7VyYlJeV2SEC+aN26tTG/8sorPWbp6enGttnlvhg3bpwxX7BggTE/duyYMb/++us9Zq+88oqxbXYeffRRYz527Fif+sfF7dVXXzXmoaGhHrObb77Z2PbkyZNejcnusvt9NLvjYF4eyy5mXHEFAAAAAAAAW6JwBQAAAAAAAFuicAUAAAAAAABbonAFAAAAAAAAW6JwBQAAAAAAAFuicAUAAAAAAABbonAFAAAAAAAAWwoq6AEga//++68xX7Rokdd9//jjj1639dXtt99uzCMiIoz5hg0bjPm0adNyPSYgPzRq1MiYFylSxOu+s9vvlyxZ4nXfAHKvdevWPrU/dOiQn0YC+FeVKlWM+RdffGHMy5Qp48fRuEtISDDmX331lcds0KBBxranTp3yakwZTGPr06ePsW1UVJQxHz58uDEvVqyYx+yDDz4wtj1z5owxh/1169bNmLdv396Yb9++3WO2Zs0ar8ZU2L3yyivGPD093ZgvXrzYmCcmJuZyRJcGrrgCAAAAAACALVG4AgAAAAAAgC1RuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC1RuAIAAAAAAIAtUbgCAAAAAACALQUV9ABw8YmOjvaYjRkzxtg2IMBcS3399deN+dGjR405kFdmzZplzG+66Sav+/7000+N+auvvup13wD8r27duj61Hz58uJ9GAvhXUJD5V4cyZcrk2bqXLFlizLt3727MDx8+7M/h5EpCQoLH7M033zS2fffdd415SEiIMTcdT+bMmWNsu2PHDmMO+7vjjjuMeXb7T3a/u12sqlSp4jG79957jW3T0tKM+eDBg435mTNnjPmliiuuAAAAAAAAYEsUrgAAAAAAAGBLFK4AAAAAAABgSxSuAAAAAAAAYEsUrgAAAAAAAGBLFK4AAAAAAABgSxSuAAAAAAAAYEtBBT0AXHwee+wxj1lUVJSx7b///mvMt23b5tWYAF+VL1/emDdv3tyYFy1a1JgfPnzYYzZ48GBj25MnTxpzAP7XrFkzj9mDDz5obPvbb78Z8x9++MGrMQGF2Zo1a4z5Qw89ZMxN51E7mzNnjjG/9957jXnjxo39ORwUMmFhYcbcdK7KibFjx/rUvrDq06ePx6xMmTLGtlu2bDHmixYt8mpMlzquuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC1RuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC0FFfQAUPi0aNHCmL/44ote9925c2djvnHjRq/7Bnzx1VdfGfPIyEif+v/ss888Zjt27PCpbwD+17ZtW49Z6dKljW3nz59vzFNSUrwaE1DQAgK8/5t406ZN/TiSwsPhcBjz7N5TX97zuLg4Y37//fd73TfyR9GiRY15hQoVjPnnn3/uz+FcNKpXr+51W35fzRtccQUAAAAAAABbonAFAAAAAAAAW6JwBQAAAAAAAFuicAUAAAAAAABbonAFAAAAAAAAW6JwBQAAAAAAAFuicAUAAAAAAABbCiroAaDwad++vTEPDg72mP3444/GtitWrPBqTIA/dOrUyWPWoEEDn/pevHixMR84cKBP/QPIX/Xq1fOYWZZlbDtjxgx/DwfIF4888ogxT09Pz6eRXDw6duxozK+55hpjnt17bsrj4uKMbWF/J06cMObr16835ldffbUxL126tMfs6NGjxrZ2Fh0dbcy7devmdd/Lli3zui0844orAAAAAAAA2BKFKwAAAAAAANgShSsAAAAAAADYEoUrAAAAAAAA2BKFKwAAAAAAANgShSsAAAAAAADYEoUrAAAAAAAA2FJQQQ8A9lO8eHFjfvPNNxvz06dPe8wGDhxobHvmzBljDvgiMjLSmL/88sses+DgYJ/WvX79emN+8uRJn/oH4F/lypUz5rGxsR6zbdu2GdvOnDnTqzEBBa1jx44FPQRbioqKMua1atXymJl+9vCHQ4cOecz4ubvwS05ONuY7duww5rfffrsx/+677zxm7777rrFtXqpTp44xr1atmjGvUqWKMbcsK7dDcklPT/e6LTzjiisAAAAAAADYEoUrAAAAAAAA2BKFKwAAAAAAANgShSsAAAAAAADYEoUrAAAAAAAA2BKFKwAAAAAAANgShSsAAAAAAADYUlBBDwD289xzzxnza665xpjPnz/fY7Z8+XKvxgT4w3/+8x9j3rhxY6/7njVrljEfOHCg130DyH89e/Y05tHR0R6zefPm+Xk0AOzslVdeMeaPPfZYnq17165dxrxHjx4es927d/t5NLCb7H7+dDgcxrxDhw4es88//9yrMfnD4cOHjbllWca8TJky/hyOm4kTJ+ZZ35cyrrgCAAAAAACALVG4AgAAAAAAgC1RuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC1RuAIAAAAAAIAtBRX0AJD/TF9rKkmvvfaaMT9+/Lgxf/3113M9JiA/PPPMM3nWd//+/Y35yZMn82zdAPwvJibG67b//vuvH0cCoKDNnTvXmF955ZX5NJLMNm/ebMyXLVuWTyOBHW3dutWY33nnnca8fv36HrMaNWp4MyS/mDFjhk/tJ02aZMzvvfder/tOTk72ui0844orAAAAAAAA2BKFKwAAAAAAANgShSsAAAAAAADYEoUrAAAAAAAA2BKFKwAAAAAAANgShSsAAAAAAADYEoUrAAAAAAAA2FJQQQ8AeSMyMtJj9v777xvbBgYGGvO5c+ca85UrVxpz4GJUunRpY37mzJl8Gklmx44dM+amsQUHBxvbhoWFeTWmDOHh4R6zZ555xqe+s5OWluYxe+GFF4xtT5065e/hwGZuvfVWr9t+8803fhwJYB8Oh8OYBwR4/zfxW265xeu2kjR+/Hhjftlll3ndd3bblZ6e7nXfvurYsWOBrRsXv/Xr13uV2d1ff/2VZ33XqVPHmG/cuDHP1n0x44orAAAAAAAA2BKFKwAAAAAAANgShSsAAAAAAADYEoUrAAAAAAAA2BKFKwAAAAAAANgShSsAAAAAAADYEoUrAAAAAAAA2FJQQQ8A3gkMDDTm8+fP95hVrVrV2HbHjh3G/LXXXjPmwKXo999/L+ghePTll18a83379nnMypYta2x71113eTUmu9u/f78xHzJkSD6NBHmlZcuWxrxcuXL5NBKg8Bg7dqwxHz58uNd9f/vtt8Y8PT3d67790b6g+h43blye9Q1cqhwOh0+5ycaNG71uC8+44goAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2FFTQA4B3qlevbswbNmzodd/PPPOMMd+xY4fXfQMFae7cucb8tttuy6eR5K877rijwNZ99uxZY56enu5133PmzDHma9as8brvpUuXet0WhUOXLl2MeWBgoDH/7bffPGY///yzV2MC7O7rr7825s8995wxj4qK8udwbOPQoUPGfMuWLR6zPn36GNvu27fPqzEB8MyyLJ9y5D+uuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC1RuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC0FFfQAkLWYmBhjvmDBAq/7zu6rir/99luv+wbsrGvXrsb8+eef95gFBwf7ezhuateu7TG766678nTdn3zyicds165dPvX91VdfGfOtW7f61D/gSUhIiDFv3769T/3PmDHDY5aWluZT34BdJSQkGPPu3bsb886dO3vMnnzySW+GZAtDhgwx5h9++GE+jQRAThQrVszrtsnJyX4cCXKKK64AAAAAAABgSxSuAAAAAAAAYEsUrgAAAAAAAGBLFK4AAAAAAABgSxSuAAAAAAAAYEsUrgAAAAAAAGBLFK4AAAAAAABgSw7LsiyfO3E4/DEWnGfIkCHG/KWXXvK67yZNmhjzNWvWeN038oYfpqkb5iyQt5izBS84ONiYL1myxJgfPHjQmN9zzz0es1OnThnbwn6YswXv5ptvNuZ9+vQx5h07djTmc+bM8ZiNHz/e2Da7z3Pz5s3GfPfu3cYcucechS/2799vzIOCgjxmb7zxhrHte++959WYLna+zlmuuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC1RuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC1RuAIAAAAAAIAtOSzLsnzuxOHwx1guKS1btjTmc+fONeYlSpTwet1NmjQx5mvWrPG6b+QNP0xTN8xZIG8xZ4HChTkLFC7MWfjim2++Mebvvvuux2zRokX+Hs4lwdc5yxVXAAAAAAAAsCUKVwAAAAAAALAlClcAAAAAAACwJQpXAAAAAAAAsCUKVwAAAAAAALAlClcAAAAAAACwJQpXAAAAAAAAsKWggh7ApSo2NtaYlyhRwqf+d+zY4TE7efKkT30DAAAAAFAYdezYsaCHgFziiisAAAAAAADYEoUrAAAAAAAA2BKFKwAAAAAAANgShSsAAAAAAADYEoUrAAAAAAAA2BKFKwAAAAAAANhSUEEPAN6Jj4835jfccIPH7OjRo/4eDgAAAAAAgN9xxRUAAAAAAABsicIVAAAAAAAAbInCFQAAAAAAAGyJwhUAAAAAAABsicIVAAAAAAAAbInCFQAAAAAAAGyJwhUAAAAAAABsyWFZluVzJw6HP8YCwAM/TFM3zFkgbzFngcKFOQsULsxZoHDxdc5yxRUAAAAAAABsicIVAAAAAAAAbInCFQAAAAAAAGyJwhUAAAAAAABsicIVAAAAAAAAbInCFQAAAAAAAGyJwhUAAAAAAABsyWFZllXQgwAAAAAAAAAu5PcrrhyDHOo/t3+2r5u4fqIcgxzalbjL30PIZFfiLjkGOfT28rfzfF35LeN9XPPPmoIeSr67mD9X5F5Ojz15uf64xXGuf/tyjFu8a7EcgxxavGux1+OpMqqKes7q6XX78/155E/dNPkmhb0VJscgh2ZtneWXfpFznFvzF+fWwvO59pzVU1VGVSnoYeSOwyH1z8H5auJE52t37crrETnX4XBIbxeOzz1XMt7HNZfefC40n2t+7usecJ7NX5xnC8/napfzrFeFq+V/L1fc4jglpiT6eTjwZMyvYzRx/cSCHgZQoC7lY8/UDVM1auWoAll3j1k9tOHgBg25fogmd5msRpc1KpBxXOwu5f27oHBuRZ5ZvlyKi5MSEwt6JJeOMWOcBRDAA86z+Y/zLPzF68LVoCWDfJr09199v5JfSVZMWIzXfVxKmPSAf449+cWXY1yrmFZKfiVZrWJauZbltnC1rf82fdzx41yv+0LJZ5K1Ys8K9bqml/o36a/7rr5PFUtV9LlfZMa5Nf9xbi38Pu74sbb131bQw8hs+XJp0CDfClf33y8lJ0sxzOccoXBVeOXTvs55Nv9xni387HKeDSqoFQcGBCowILCgVg/kWNLpJIUWCS3oYaCQ8eUYF+AIULGgYj6tv2hQUZ/aZzh06pAkKbxYuF/6yw3LspRyNkXFg4vn+7oLK86tKCz8dW4NDgz2w2hsKjDQ+R9gd0lJUqgP87kQ7eucZ1FYXGzn2VxfcRW3OE7P/fCcJKnqe1XlGOTI8j7fWVtnqc6YOio6uKhqj6mt+dvnu+VZ3R+85p81avdZO5UZXkbFhxRX1feq6qHZD2U7pty0G792vKq/X11FBxdV448b69e9v2Z6zU87f1LshFiFDg1V+Fvhuu2L27Tl0BZX/vuB3+UY5NCcbXNcy9b+s1aOQQ41+KiBW1+3TLlFTf/b1KdtrDKqijYd2qQlCUtc73ebiW3cXpN6NlXPfP+MokZEKXRoqLpM66JDSYcy9TXvz3mubSv5Zkl1mNpBmw5uMq5fOvd5/bL7l2zXc+Gzfs7fjvOfuZPR57Ldy/TEvCcUNSJK4W+Fq+83fXU67bQSUxL1wMwHFDEsQhHDIvT8D8/L03cJjFwxUjGjYlR8SHG1nthaGw9uzPSarYe3qtv0bio9rLSKDS6mRuMbuX2G549pya4l6vddP0WPiFbFkearSw4mHVSv2b1U9u2yKja4mOqNq6dJ6ye5veb8e5lzsg/mZKyXGn8deyRp7/G9emj2Qyr7dlnX6z757ZMcjSP1bKqenv+0okZEqeSbJdXp807ac3xPptdldYxLt9IVtzhOl71zmUKGhOi6Sddp86HNmebGhc+4ajOxjb778zslHEtwbXd295pf2OeZtDMatHiQLh99uYoNLqbI4ZFq+UlL/bDjB499xC2OU8wo518Un/vhuUzr/W3fb7plyi0q9WYplRhaQjd8eoNW7lmZqQ/HIEeO3p8qo6ro1qm36vvt36vR+EYqPqS4Plr7kXE7LxacWzm3cm49J+MYOH3TdA35eYgqvltRxQYX0w2f3qDtR7e7vfbCZ2/Y4nwbFyc955zPqlrV+eyerJ7fM2uWVKeOVLSoVLu2NP+C81VWz/1Zs0Zq104qU0YqXtzZ/0PZz+dctRs/Xqpe3Tmuxo2lXzO/b/rpJyk21lmsCA+XbrtN2nJuPuv3351jn3Pee7l2rXNZA/f5rFtukZqem89ebWOVKtKmTdKSJefe7zZt3F+Tmio984wUFeUcd5cu0qHM81nz5p3btpIlpQ4dnH1nJ+Pz+uWX7NfjcDj3k6y2o2fPzH0uWyY98YSzz/BwqW9f6fRp5xV9DzwgRUQ4/3v+ecnT926NHOm8oql4cal1a2lj5vmsrVulbt2k0qWlYsWkRo3cP8Pzx7RkidSvnxQdLVXM5krs0aOd+3hIiHOcjRpJU6dm7jNjX//pJykgQBowwL2fqVOdrxs71ry+LHCe5TzLefacwniezfUVV11rdtUfR/7Q5xs/18h2I1UmpIwkKSokyvWaZbuX6estX6tf434qWaSk3l/9vm6ffrt2P7VbkSGRWfZ7MOmgbpp8k6JCo/RiyxcVXixcuxJ36estXxvHk5t2UzdM1YnTJ9S3YV855NDw5cPVdXpX/fXEX65K4sK/FuqWKbeoWkQ1xbWOU/LZZI1ePVotPmmhdX3XqUp4FdWJrqPwYuH6OeFndbqykyRp6e6lCnAEKP5AvI6nHlepoqWUbqVr+d/L1adBH5+2cdTNo/T4vMdVokgJvRL7iiSpbGhZt9c8Pu9xRRSP0MDWA7UrcZdGrRyl/oH9Na3bNNdrJsdPVo9ZPdSuRjsNaztMp86c0tg1Y9VyQkv91vc3VQmvYhxHTteTW4/Pe1zlSpTToDaDtHLPSo1fN17hxcK1fM9yVQ6rrKHXD9Xc7XM1YvkI1YmuowfqPeDW/tP4T3Xi9Ak91vgxpZxN0Xur3tP1k67Xhkc3qGwJ5/u06eAmtfikhSqUqqAXW76o0OBQTd88XZ2/6Kyv7vxKXWp2ceuz39x+igqJ0oDWA5R0Osnj2JPPJKvNxDbafnS7+jfpr6rhVfXl5i/Vc3ZPJaYk6slmT7q9Pif7YG7Heqnw17HnwMkDava/ZnLIof6N+ysqNErzts9Trzm9dDz1uJ5q9pRxHA9/87A++/0z3VP3HjWv2Fw/7fpJHaZ2yNE2vLTwJQ1fPlwdr+iodtXbKf5AvNp91k4pZ1OM7V6JfUXHUo9pz/E9GtlupCSpRJESOVpnhrjFcXpz2Zt6uMHDalKhiY6nHteaf9Zo3b51urH6jVm26Vqzq8KLhevp75/W3XXuVvvL27vWu+ngJsVOiFWpoqX0fIvnFRwQrI/WfqQ2E9toSc8lalqxaZZ9ZmfbkW26+6u71bdhX/Vu0FtXlrnSq34KG86tnFs5t2b21rK3FOAI0LPNn9WxlGMavny47v36Xq16eFW2bQv0fNu1q/THH9LnnzuLBWWc81lR5+azli2Tvv7a+Yt/yZLS++9Lt98u7d4tRWY9n3XwoHTTTc5+XnzRWcDYtcvZj0lu2k2dKp044SyMOBzS8OHO7fnrLyn4///qvnChs9hUrZqz+JKc7CxMtGghrVvnLL7UqeNcz88/S52c81lLlzqLEfHx0vHjUqlSUnq687bKPn1828ZRo6THH5dKlJBecc5nlXWfz3r8cWfRZOBAZ5+jRjkfkj/tvHk2ebLUo4ezcDZsmHTqlLNI0rKl9Ntvzm3LTk7Wk1uPPy6VK+e8/XTlSmdxMTzc+d5VriwNHSrNnSuNGOF87x9wn8/69FPn5/rYY1JKivTee9L110sbNpx7nzZtcn6GFSo43/vQUGn6dKlzZ+mrr5wFuPP16+f8nAYMcF5x5cnHHzuLbt26SU8+6Vz/779Lq1ZJ99yTdZvrr3f2/+abzvU3aCDt2+d8H9q2lR55JNdvIedZzrOcZzMrVOdZywsjfhlhKU7Wzn93ZsoUJ6vIG0Ws7Ue2u5bF74+3FCdr9KrRrmUTfpvg1sfMLTMtxcn6de+vuRpLTtrt/HenpThZkcMiraOnjrqWz94621KcrG+2feNaVn9cfSt6RLR15NQRt/EHDAqwHpj5gGtZhykdrCYfN3H9u+u0rlbXaV2twEGB1rw/51mWZVnr/llnKU7W7K2zfdpGy7Ks2h/WtlpPaJ1pecb72PbTtlZ6erpr+dPzn7YCBwVaicmJlmVZ1onUE1b4W+FW7zm93drvP7HfCnszLNNyb9djWc59YOCigZn6iBkZY/WY2SNTn+0mt3Pr89r/Xms54hzWI9884lp2Nu2sVfHdim7vQcbnWnxwcWvPsT2u5av2rLIUJ+vp+U+7lt0w6Qar7pi6VsqZFNey9PR0q/n/mluXv395pjG1/KSldTbtrPE9sSzLGrVilKU4WZ/Ff+Zadvrsaeva/15rlRhawjqectxtrDnZB3M61kuRP449vWb3ssq/Xd46nHTYrX33Gd2tsDfDrFOnT3lc//p96y3Fyer3bT+35fd8dU+m/f7CY9z+E/utoNeDrM5fdHZrG7cozlKc3ObGop2LLMXJWrRzkWtZhykdrJiRMR7HdqEL51u9sfWsDlM65Lh9hox9d8QvI9yWd/6is1XkjSLWjqM7XMv+Of6PVXJoSavVhFauZQMXDbQUl/lUc+H7kzFmxcma/+f8XI/zYsC5lXOrp/VY1qV1bs04Btb8oKaVejbVtfy9le9ZipO14cAG17IeM3u4HRttc74dMcKyJMvauTNzJllWkSKWtf3cfLbi453LR5+bz9aECe59zJzp/PevudzXc9Ju507nayIjLevouffNmj3bufybc++bVb++ZUVHW9aRc/PZio+3rIAAy3rg3Hy2OnSwrCbn5rPVtavzv8BAy5rnnM/WunXO/mfP9m0bLcuyate2rNatMy/PeB/btrWs8+aE9fTTzrEk/v88O3HCssLDLav3BfN2/37LCgvLvNzb9ViW83UDB2buIybGsnr0yNxnu3bufV57rWU5HJb1yLn5bJ09a1kVK7q/Bxmfa/HilrXn3Hy2Vq1yLn/63Hy2brjBsurWtayUc/PBSk+3rObNLevy8+ZDxphatnSuMzu33eb8bEwu3Ncty7KSkiyrRg1n25QU5/5UqpRlJSRkv04POM9ynvW0HsviPGtZ9j7PBuS8xJVzbau1VfXS1V3/vrrs1SpVtJT++vcvj20ynp/y7R/f6kzamRyvKzft7qp9lyKKR7j+HVs5VpJc49p3Yp/W71+vnvV6qnTx0m7jv7HajZr751y3tuv2rXNVMpftXqb2Ndqrfrn6WpqwVJKzgu2QQy0rt/RpG3OiT4M+cjjO3Y4TWzlWaVaaEo4lSJJ+2PGDElMSdXedu3X41GHXf4EBgWpasakW7Vrkl/V4o9c1vdz6bFqhqSxZ6tWgl2tZYECgGl3WKMt9qPNVnVWhVAXXv5tUaKKmFZq6Pq+jyUf1086fdGftO3Xi9AnXth9JPqJ21dvpz6N/au/xvW599m7QO0f3r8/dPlflSpTT3XXvdi0LDgzWE02f0MnTJ7UkYYnb67PbB70ZK87J7thjWZa+2vKVOl7RUZYst7nQrno7HUs9pnX71nnsP2OfeqLpE27Ln2r6VLZj+3Hnjzqbflb9GvVzW/5408dzunk+CS8Wrk2HNunPI3/63FdaepoW7Figzld1VrWIaq7l5UuW1z1179Gy3ct0PPW4V31XDa+qdjXa+TzGixHnVs6tuVGYz60ZHqz/oIoEFnH9+8J9y8T259u2bZ2342W4+mrnFUh/GbYtPNz5v99+K53Jxb6em3Z33eW8WihDrPN9c41r3z5p/Xrn7Wylz81nXX21dOONzqt+zm+7bt25q3GWLZPat5fq13defSU5/9fhcF7R5Ms25kSfPs51nT++tDQp4f/n2Q8/OG+9u/tu6fDhc/8FBjpvZVyUs/mc7Xq80auXe59NmzpvCex1bj4rMNB5C15W+1Dnzs4rqTI0aeLsI+PzOnrUeXvenXc6r8zK2PYjR5xXn/35p7T3gvnQu3fOnksVHi7t2ZP1LacmISHOWwi3bJFatZK++855BWPlyrnrJxc4z3KezQ3Os/l7ns2Th7NXLpX5gBJRLEL/pvzrsU3rmNa6vebtGrRkkEauHKk2Vdqo85WddU/de4wPGc5Nu8ph7uPKeKP/TXaOK2PHzerWlJplaur7Hd+7HnIWGxOrs+lntWLPClUqVUkHkw4qNiZWmw5t0tLd5yZ9rahargOIt9uYE9lt259Hnb+sXv/p9Vm2L1W0lF/W440L+wwrFiZJqlSqkvvyomFZ7kOXl74807IrIq/Q9E3TJUnbj26XJUuvLXpNry16LcsxHEw66HbgqBpeNUdjT0hM0OWlL1eAw70GXLNMTVd+vuzeP2/GinOyO/YcOnVIiSmJGr9uvMavG59lHweTDnrsP+FYggIcAW4/1EhZHzMytf3/faFG6Rpuy0sXL62IYhFZNfGr1697Xbd9cZuu+OAK1Ymuo5ur36z7692vq8teneu+Dp06pFNnTunKyKyPlelWuv4+9rdqR9fOdd9VI3I29y5FnFs5t+ZGYT63etqG3Lwvtj/fZvXLd0SE9K9h21q3dt5OOGiQ8xf4Nm2cBYl77nE+j8of7S4cV0YRK2NcGcWXK7M479WsKX3//bkHdcfGSmfPSitWSJUqOW8DjI113pJ2fuGqVq1zRTBvtzEnstu2P///DzvXZz2fVSpn8znb9Xjjwj7DnPNZlSplXp7Vei7PPJ91xRXOWwElaft2ZyHstdec/2Xl4EH34lfVHM7nF15w3l7apIlUo4bzVtB77nHelpidFi2kRx+VPvzQWUDLyfPcfMB5lvNsbnCezd/zbJ4UrjxV+SxPDwuU5HA4NOPOGVq5Z6W+2faNvt/xvR6a85DeWfGOVj680uPzXHLTzuO45HlcnjS6rJGKBRXTzwk/q3JYZUWHRuuKyCsUWzlWY34do9SzqVqasFRdrjp336a325gT2W1bupUuSZrcZbLKlSiX6XVBATnbFXx5D9OstFz1mdVy0z7kSca2P3vtsx6v5LiwmJBX32KW088pN2PFOdkdezLe3/uuvk896vXI8rXeFHIKg1YxrbTjiR2avXW2Fvy1QP/97b8auXKkxt06Tg83eDjP1utQ5gezS86rtrJSPIhvEPSEcyvn1qxczOdWX94X259vPV2pYvosHA5pxgznM46++cZZJHroIemdd5zLSnjY13PTzptxedKokfMB3z//7Cy8REc7iyWxsdKYMc6HpS9d6v7sJG+3MSey27Z05z6hyZOdz5O6UFAOf23y5T1My3o+e+wzq+XefFYZ2/7ss84CUVZqXDAfiudwPtesKW3b5ryKbv585/OyxoxxPhtr0CBz29RUafFi5//fscP5zLGQkJyt1wucZznPZoXzrHdt/X2e9apw5ekXEX9oVrGZmlVspiE3DNHUDVN179f36ouNX2T7i5W37c4XE+b89qxth7dlyrYe2aoyIWVcXylZJLCImlRooqW7l6pyWGXXpXGxMbFKTUvVlA1TdCDpgFrFtPLLWM+/DNEb1SOcV4hEh0arbbW2PvWVnYhiEUpMSXRbdjrttPad2Jcn68uoxJ/vjyN/uB7Ul3ErU3BgsN+3PSY8Rr8f+F3pVrrbVVdbD2915bmRl2O9GPh67IkKiVLJIiWVlp7m1fsbExajdCtdO47ucPurVlbHjExt/39f2H50u9tVRUdOHTH+JS+Dr8cAyXl114PXPKgHr3lQJ0+fVKsJrRS3OC7XhauokCiFBIdo25EsjpWHtyrAEaBKYc6/NmX89SUxJdF1qbkkny7NvlhxbuXcanIpnVvzQ56P3w/HbI+aNXP+N2SI82Hq994rffGF9HA289LbdueL+f+fa7Zlcd7butX5IPrQ///69SJFnFfZLF3qLFxl3HYYG+ssSEyZIh044LwNzB9j9fU9z7h1MzraeStnXoqIcN6WeL7Tp523YuaFPzPPZ/3xx7mHzVf7/9v+g4PzZttDQ523od51l3M7u3Z1frYvveQsbnoycKDzVsG333ZeufXii84vMvAS51nOsyacZ/3L3+P36hlXGTv+hR+sL/5N/jdTJbJ+ufqSnF+T6e92WSlfsrzql6uvSfGT3LZt48GNWrBjgdpf3t7t9bGVY7Vqzyot2rnINenLhJRRzTI1NeyXYc7XxMT6ZayhwaE+vd/tarRTqaKlNHTp0CzvTc7qa0e9Vb10df2c8LPbsvFrx3usVvtq1tZZbvfHrt67Wqv2rtItNW6R5DzQtanSRh+t/SjLA48v296+RnvtP7lf0zae+0aKs+lnNXr1aJUoUkKtY1rnqr+8HOvFwNdjT2BAoG6vdbu+2vJVll83m937e8vlzn3q/VXuPzSNWjUq23XfUPUGBQUEaewa969w/mD1B9m2lZzHgGOpx3L02qwcOXXE7d8lipRQjdI1lJqWu+Ok5Hwfb6p+k2Zvne32ddAHTh7Q1I1T1bJyS9el2xk/cJx/TEg6naRJ8ZO82IqLG+dWJ86tWbuUzq35Ic/Hn1G8ubA44Yt//818NU39+s7/TTXs6962y0r58s62kya5b9vGjdKCBc5nWJ0vNtb57XGLFp0rXJUp47wKZ9iwc6/xx1hDQ317v9u1c94OOHRo1s/XOuTHfbp6deeVaOcbP97zFVe+mjXL/RlVq1c7P5dbnPNZ0dHO2zI/+ijr4pkv237E/ecPFSnivD3UsszPMVu1ylmweuop6T//kZ57TvrgA2nJEs9tssF51onzbNY4z/qXv8fv1RVXDcs3lCS98tMr6l67u4IDg9Xxio6ug4E3JsVP0phfx6jLVV1UvXR1nUg9oY/XfaxSRUtlmmz+aOfJiBtH6JYpt+ja/12rXtf0UvIZ51eJhhUNU1zrOLfXxlaO1ZClQ/T38b/dJnermFb6aO1HqhJeRRVLVfTLWBuWb6ixa8Zq8M+DVaN0DUWHRuv6qh7uwc9CqaKlNLbDWN0/8341GN9A3Wt3V1RolHYf263v/vxOLSq10Aftc/YLdHYevuZhPfLdI7p9+u26sdqNit8fr+93fO/62ll/q1G6hlpOaKlHGz2q1LOpGrVqlCKLR+r5Fs+7XvNh+w/V8pOWqju2rno36K1qEdV0IOmAVuxZoT3H9yj+kXiv1t2nYR99tPYj9ZzdU2v3rVWV8CqasXmGfvn7F41qN0oli5bMdZ95NdaLgT+OPW/d8JYW7Vykpv9tqt4NeqtWVC0dTT6qdfvWaeFfC3X0haMe29YvV19317lbY9aM0bHUY2peqbl+3Pmjth/dnu16y5YoqyebPql3VryjTp930s01blb8/njN2z5PZULKZPsXqYblG2rapml65vtn1PiyxipRpIQ6Xtkxx9tda0wttanSRg3LN1Tp4qW15p81mrF5hvo36Z/jPs43+LrB+mHHD2r5SUv1a9xPQQFB+mjtR0o9m6rhbYe7XndT9ZtUOayyes3ppeeaP6dAR6A+Wf+JokKcxx+cw7nViXNr1i6lc2t+ydPxN3TOZ73yitS9u/NKlo4dzxW0vDFpkvMWqy5dnIWPEyekjz92FlsuLBj5o50nI0Y4Cx7XXut8OHhysjR6tPP5SnFx7q+NjXVeWfP33+4FqlatnEWSKlWkiufms09jbdhQGjtWGjzYeVtbdLTn51VlpVQpZ/v775caNHB+blFR0u7dzgeDt2jhLJz4w8MPS4884nye1403SvHxztsiy+TNfFaNGs4H4D/6qLMAOGqUFBkpPX9uPuvDD52vqVvX+eD1atWcV8StWOF8uHq8l/Phppuct162aCGVLeu8guqDD6QOHaSSHn5OTkmRevRwPptryBDnskGDnLePPvigtGGDV3OJ86wT59mscZ71P3+O36vCVeMKjfXGdW9o3Jpxmr99vtKtdO18cqdPk751TGut3rtaX2z6QgdOHlBYsTA1qdBEU7pOMT6s19t2nrSt1lbz752vgYsHasCiAQoODFbrmNYa1nZYpv6aV2quQEegQoJDVK9sPdfy2Mqx+mjtR64Ktj/GOqD1ACUcS9DwX4brxOkTah3TOleTXpLuqXuPLit5md5a9pZGLB+h1LRUVShZQbExsXqw/oO56sukd8Pe2pm4U//77X+av32+YivH6of7f9ANn97gt3Wc74F6DyjAEaBRK0fpYNJBNanQRB+0/0DlS5Z3vaZWVC2t6bNGg5YM0sT4iTpy6oiiQ6N1TflrNKDVAK/XXTy4uBb3XKwXF76oSfGTdDz1uK6MvFITbpugnvV7etVnXo31YuCPY0/ZEmW1uvdqvb7kdX295WuN+XWMIkMiVTuqtoa1HZZt+09ucxZdpmyYollbZ+n6qtfru3u+U6WRlbJtO6ztMIUEh+jjdR9r4V8LdW2la7Xg/gVq+UlLFQs0XCovqV/jflp/YL0mrJ+gkStHKiYsJleFqyeaPKE5f8zRgh0LlHo2VTHhMRp8/WA91/y5HPdxvtrRtbX0waV66ceX9OayN5Vupatphab6rMtnalqxqet1wYHBmnnXTPX7rp9eW/SaypUop6eaPqWI4hF6cLb/jjsXA86tTpxbs3YpnVvzS56Ov3Fj6Y03pHHjnM/1SU+Xdu70rXDVurXzSpkvvnAWFMLCnLfiTZliflC2t+08advWuU0DBzqfUxQc7FzHsGGZ+2ve3PkcppAQqd65+azYWGfhKtZ9Pvs01gEDnA+PHz7cWfBq3Tp3hSvJ+dDwyy6T3nrLWaBLTXU+kDw21lkw8ZfevZ37w//+53wvY2Od32p4Q97MZz3wgBQQ4CxYHTzofE8/+MB5BV2GWrWkNWucBaKJE51XSkVHS9dc43xvvdW3r/Pze/dd6eRJZ6HyiSekV1/13Obll50PjF++/NythEWKOAubzZo5r74aMybXQ+E868R5NmucZ/3Pn+N3WN48KQwA4BeJKYmKGBahwdcN1iutXino4QAAAACArXj1jCsAQO4ln0nOtGzUylGSpDZV2uTvYAAAAACgEPDqVkEAQO5N2zRNE9dPVPvL26tEkRJatnuZPt/4uW6qfpNaVG5R0MMDAAAAANuhcAUA+eTqslcrKCBIw38ZruOpx10PbB98/eCCHhoAAAAA2BLPuAIAAAAAAIAt8YwrAAAAAAAA2BKFKwAAAAAAANgShSsAAAAAAADYkl8ezu5wOPzRDQAP/P0oOuYskLeYs0DhwpwFChfmLFC4+DpnueIKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALYUVNADAAAAAICLVUREhMescuXKebruhIQEj9nTTz9tbLtx40Zj/scffxjz+Ph4Yw4AOcUVVwAAAAAAALAlClcAAAAAAACwJQpXAAAAAAAAsCUKVwAAAAAAALAlClcAAAAAAACwJQpXAAAAAAAAsCUKVwAAAAAAALCloIIeAAqfjh07GvM5c+Z4zPr3729sO27cOGOelpZmzAFfREdHe8ymT59ubLt8+XJjPn78eGO+a9cuY36xCgsL85i1atXK2Hb+/PnG/MyZM16NCQCA83Xo0MGYd+rUyZi3adPGY1ajRg1vhpRjf/zxh8csJibG2LZo0aI+rTswMNCn9gCQgSuuAAAAAAAAYEsUrgAAAAAAAGBLFK4AAAAAAABgSxSuAAAAAAAAYEsUrgAAAAAAAGBLFK4AAAAAAABgSw7LsiyfO3E4/DEW2ERkZKQxX79+vTGvWLGi1+sOCQkx5snJyV73XZj5YZq6uVTnbEREhDE3fWV0WFiYse3MmTON+V133WXML1bZvW9r1671mEVFRRnbNmzY0Jhv377dmOcl5qz9lSpVypi/+eabxrxOnToes7Zt2xrbnjlzxpgj/zFnC7/q1asb88cee8xj1rt3b2Pb4sWLG3M+76wFBgbmWd/MWaBw8XXOcsUVAAAAAAAAbInCFQAAAAAAAGyJwhUAAAAAAABsicIVAAAAAAAAbInCFQAAAAAAAGyJwhUAAAAAAABsicIVAAAAAAAAbCmooAcA+2nVqpUxr1ixotd9f/7558Y8JSXF676BMmXKGPNp06YZ89KlS3vMxowZY2z7+OOPG/NL1auvvmrMq1at6jHr27evse327du9GhMuDffee68xHzJkiDGvVKmS1+suVaqUMT9y5IjXfQPIWnY/nz755JP5NJL8tXXrVmO+adOmfBoJ4F81atTwmGX3M3+XLl2MeZs2bYx5enq6x2zcuHHGtr/88osx5+dX73DFFQAAAAAAAGyJwhUAAAAAAABsicIVAAAAAAAAbInCFQAAAAAAAGyJwhUAAAAAAABsicIVAAAAAAAAbInCFQAAAAAAAGzJYVmW5XMnDoc/xoJ8UrRoUWP+yy+/GPOGDRt6ve727dsb83nz5nnd98XMD9PUzcU6Z2+66SZj7sv+Va5cOWN+6NAhr/suzGrXrm3MN2zYYMxnzpzpMevZs6ex7YkTJ4x5QWLO5o+KFSt6zH777Tdj28jISGPuy2c4bdo0Y96/f39jfvToUa/XDe8wZ/2jTJkyxvzJJ5/0mGX38+f8+fONebNmzYz53LlzPWZJSUnGtqGhocZ8wYIFxnzjxo0es1WrVhnbZncsS05ONubZbVthxZy1vzp16hjz7M6FXbt29Zhld6wpSGfPnjXm27ZtM+bLli3zmJmOoZJ0+vRpY16QfJ2zXHEFAAAAAAAAW6JwBQAAAAAAAFuicAUAAAAAAABbonAFAAAAAAAAW6JwBQAAAAAAAFuicAUAAAAAAABbonAFAAAAAAAAWwoq6AEg/9WtW9eYN2zY0Kf+z5496zGbN2+eT33j0hYdHW3Mb7/9dp/679Wrl8fs0KFDPvVdWNWuXduYL1y40Kf+Z86c6TE7ceKET33j4vfss896zEqXLp2PI3F31113GfObb77ZmA8ZMsSYjx492mN2+vRpY1vAF6GhocZ8wYIFxrxevXoesy5dung1pgwrV6405g0aNPCY7dq1y9i2cuXKxnzPnj3GPD093ZgDdnX11Vd7zB577DFj2+zOhaVKlfJqTJK0d+9eY7506VJjvnPnTmP+/PPPe8zWrl1rbNukSRNjnt3PJ+3bt/eYxcfHG9uOGzfOmBdmXHEFAAAAAAAAW6JwBQAAAAAAAFuicAUAAAAAAABbonAFAAAAAAAAW6JwBQAAAAAAAFuicAUAAAAAAABbCiroASD/3X777Xnaf3ZfhQx465133jHm9913nzHP7utrv/zyy1yP6WIXGxtrzMuWLWvMJ06caMw/++yz3A4Jl5CYmBhj/uCDD3rd9++//27MDxw4YMzbtm3r9brDwsKM+bPPPmvMp0yZ4jHbv3+/V2MCJKlIkSLGfOrUqca8Xr16xnzo0KEes4ULFxrb+mrXrl1et929e7f/BgLYyEcffWTMu3Tp4jErU6aMT+v+8ccfjfmGDRs8Zi+//LKxbUpKildjytC8eXOP2aOPPmps+8knnxjz+vXrG3PTzx8ffvihse1XX31lzA8dOmTM7YwrrgAAAAAAAGBLFK4AAAAAAABgSxSuAAAAAAAAYEsUrgAAAAAAAGBLFK4AAAAAAABgSxSuAAAAAAAAYEsUrgAAAAAAAGBLQQU9AOS/Vq1a+dT+9OnTxvyVV17xqX/AE8uyjHl6erox/+eff4x5dvt2YVW8eHFj/vLLL3vM+vXrZ2yb3Wfy0EMPGXPApH79+sa8ZMmSHrOlS5ca27Zu3dqYFytWzJjffffdHjPTnJKk6tWrG/Ny5coZ89mzZ3vMbrnlFmPbo0ePGnNc/EqUKOExe+mll4xtb731VmN++PBhY/722297zE6dOmVsCyCz7M5Vzz//vDF/+OGHjbnD4fCYHTp0yNh27NixxnzEiBHGPCkpyZjnpcjISI9ZYGCgsW1cXJwxnz9/vjGPiYkx5pcqrrgCAAAAAACALVG4AgAAAAAAgC1RuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC1RuAIAAAAAAIAtUbgCAAAAAACALQUV9ACQN5o3b+5VlhNJSUnGfP369T71D+SVDh06GPMFCxZ4zBITE41tx44d682Q/KJ169bGvE2bNsa8WbNmXq97xowZXrcFslO0aFFjblmWx2zkyJE+rTslJcWYT5gwwWN2xx13GNtWq1bNqzFlOHXqlMfs9OnTPvWNi1/nzp09Zi+++KKx7e7du415bGysMT927JgxB5A72f2M99xzzxlzh8NhzPfu3esxu/32241tV69ebczzUmBgoDGvVKmSMf/00089ZnPnzjW2jYiIMObZMX0mkydPNrbN7veVwowrrgAAAAAAAGBLFK4AAAAAAABgSxSuAAAAAAAAYEsUrgAAAAAAAGBLFK4AAAAAAABgSxSuAAAAAAAAYEsUrgAAAAAAAGBLQQU9AOSNxo0b51nfY8eOzbO+AZP33nvPmF933XXG/LLLLjPmrVq18pg5HA5j206dOhnzvJTd2CzL8rrvv/76y5i//PLLXvcNZOfuu+/2um2HDh2M+axZs7zuOzuNGjXKs74laeXKlR6zkydP5um6Ufg1b97c67a//fabMd+zZ4/XfQPIvcDAQGOelpbmU/9nz571mDVt2tTYtlu3bsb8qquu8mpMkpScnGzMa9as6VN++PBhj1nZsmWNbX114MABj9ngwYONbc+cOePv4dgGV1wBAAAAAADAlihcAQAAAAAAwJYoXAEAAAAAAMCWKFwBAAAAAADAlihcAQAAAAAAwJYoXAEAAAAAAMCWHJYv35Oe0Uk2X8WO/Dd58mSP2X333Wdsm5iYaMzr1q1rzPkqZP/zwzR1c7HO2YiICGNev359Y37zzTd7zJ577jlj24MHDxrzSZMmGXNfmOa7JMXHx3vd92effWbMe/To4XXfFzPmrH/ceeedxvzzzz/3mG3YsMHYtnv37sY8u3Ndly5dPGZ33HGHse3x48eNeXbHsqNHj3rMWrVqZWy7efNmY36pupTmrOl8FRkZaWybmppqzIcNG2bMZ8+e7TFbv369sS1wvktpzpoUL17cmE+dOtWYt23b1piHhIR4zAICzNfA+PoZpaWlecwCAwN96jsvpaenG/OZM2ca8yeeeMJjtm/fPq/GZAe+7g9ccQUAAAAAAABbonAFAAAAAAAAW6JwBQAAAAAAAFuicAUAAAAAAABbonAFAAAAAAAAW6JwBQAAAAAAAFuicAUAAAAAAABbcliWZfncicPhj7EgF1q2bGnMlyxZ4jELCDDXKxMSEox5lSpVjDn8zw/T1A1z9uJSrVo1Y759+3Zjvn79eo9Zu3btjG0PHTpkzC9VzFn/KF26tDE37dthYWHGttm9p758hgsXLjTmjz32mDH/9ttvjfnll1/uMfv444+NbR955BFjfqm6lOasaVvT09PzdN2m/seNG2dsu3LlSmNeuXJlY246XmzatMnYNju1a9c25itWrPCY7dmzx6d1X6oupTmbl8LDw435iy++6DFr0aKFse2RI0eM+e7du4150aJFPWb16tUztm3SpIkxz0vZHctefvllY56YmOjH0diHr3OWK64AAAAAAABgSxSuAAAAAAAAYEsUrgAAAAAAAGBLFK4AAAAAAABgSxSuAAAAAAAAYEsUrgAAAAAAAGBLFK4AAAAAAABgS0EFPQB4JzIy0pgHBHhfk/zhhx+8bgsg/w0YMMCYW5ZlzF944QWP2aFDh7waE+APR48eNeZ33nmnx2zGjBnGtmFhYV6NKcPo0aM9ZqY5JUkpKSnG/OuvvzbmL774osesXbt2xrbVq1c35jt27DDmKPzefvttj9kzzzyTp+s2/Xzar18/Y9vscjsznUsXL15sbNu9e3c/jwY4JzEx0ZibzjcF6dNPPzXmTZo08an/EydOeMyyO05OnDjRmKelpXkzpEseV1wBAAAAAADAlihcAQAAAAAAwJYoXAEAAAAAAMCWKFwBAAAAAADAlihcAQAAAAAAwJYoXAEAAAAAAMCWKFwBAAAAAADAlhyWZVk+d+Jw+GMsyIXJkycb8/vuu89jlpiYaGx74403GvM1a9YYc/ifH6apG+Zs4XLHHXcY82nTphnzEydOGPPrrrvOY7Zu3TpjW2SNOVvw2rZta8zvueceY57duXLAgAEes5MnTxrbZqd48eLGfOrUqR6zTp06Gdt+9tlnxrxHjx7G/GJ1Kc3ZwMBAj9k111xjbGva9yQpKCjImFeqVMljFhBwaf49Pbt9Ly4uzpgPHjzYj6MpPC6lOXupev755z1m2e332R2LsnPvvfd6zD7//HOf+r5U+TpnL80zBAAAAAAAAGyPwhUAAAAAAABsicIVAAAAAAAAbInCFQAAAAAAAGyJwhUAAAAAAABsicIVAAAAAAAAbMlh+eG7RPn6UP+rWLGiMU9ISDDmpq8U3rhxo7Ft3bp1jTnyH1/5e2n75JNPjHnPnj2NeXZf22v6yl94hzmLvNS9e3eP2ZQpU4xt9+7da8zr169vzI8ePWrMCyvmbP644YYbPGbBwcHGtnFxcca8cePG3gzJ9ubMmWPMu3Tpkk8jsRfmbOH38MMPG/N3333XY1aiRAmf1r1p0yZj3qhRI49ZamqqT+u+VPk6Z7niCgAAAAAAALZE4QoAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALYUVNADQNaaN29uzAMCvK85zpo1y+u2APLfLbfcYsyTkpKM+TvvvOPP4QAoYNOnT/eYderUydj2rrvuMub9+/c35q+//roxB0x+/PFHr9vWr1/fmDdu3NiYnz171mM2YcIEY9uPP/7YmD/11FPG/J577jHmwMWoSZMmxjy7n09LlCjh9bpPnjxpzB955BFjnpqa6vW6kTe44goAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2FFTQA0DWIiMjfWp/+PBhj9l7773nU98A/O+RRx7xmJUtW9bY9uDBg8Z83bp1Xo0JgD2lp6d7zIYPH25se9tttxnzgQMHGvMvvvjCY/bHH38Y2wK+WLBggTEfMmSIMQ8K8vxrT+/evY1ta9SoYczbtGljzH2xZ8+ePOsbyEsdO3Y05iVLlvS676SkJGPeqVMnY/7LL794vW4UDK64AgAAAAAAgC1RuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC1RuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC0FFfQAkLV27dr51H737t0es2PHjvnUNwD/e+SRRzxmlmUZ23733Xc+rbtkyZIes4iICGNb07EGQP5bv369MR8wYIAxHzFihDEfOnSox+z+++83tk1OTjbmgMmWLVuM+fTp0435nXfe6fW6r7vuOq/bSlJaWprHLLtz+IsvvujTuoG8Yvr5UZKef/75PFv3lClTjPnixYvzbN0oGFxxBQAAAAAAAFuicAUAAAAAAABbonAFAAAAAAAAW6JwBQAAAAAAAFuicAUAAAAAAABbonAFAAAAAAAAWwoq6AFcqoKDg4159erVfeo/JSXFY3bmzBmf+gZgL6av2Zake++915g//fTTHrNNmzYZ2/bo0cOYA7CXTz/91Jj37dvXmHft2tVj9vrrrxvb/v7778YcMElOTjbmTz31lDEvUaKEx6xRo0bGttHR0cZ8165dxnzy5Mkes7i4OGNboCCZ5s3mzZuNbbP7fTc7pnNGdvMdFx+uuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC1RuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC1RuAIAAAAAAIAtBRX0AC5V6enpxnzNmjXGvE6dOsZ8+/btuR4TgMLp4YcfNua9evUy5v/73/88Zm+88YZXYwJgT4cOHTLmbdu2Nea7du3ymL3wwgvGtvfee68xB3xx4MABY96xY0eP2f33329s26xZM2M+aNAgY37w4EFjDtjV9ddf7zGrWLGisa1lWT6t++mnn/aYpaSk+NQ3Ch+uuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC1RuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC1RuAIAAAAAAIAtOSzLsnzuxOHwx1hwnssuu8yYDx482JivXbvWY/bhhx96NSYUHD9MUzfMWftp2bKlx+z11183tv3555+N+dixY435v//+6zE7ffq0sS2yxpzFxWrBggUes2uvvdbYtmnTpsZ88+bNXo3JH5izQOHCnM0f8fHxHrO6dev61PeIESOM+QsvvOBT/7AXX+csV1wBAAAAAADAlihcAQAAAAAAwJYoXAEAAAAAAMCWKFwBAAAAAADAlihcAQAAAAAAwJYoXAEAAAAAAMCWKFwBAAAAAADAlhyWZVk+d+Jw+GMsADzwwzR1w5wF8hZzFherUqVKeczi4+ONbZ988kljPmfOHK/G5A/MWaBwYc7mj7///ttjVrFiRWPbgwcPGvP69esb83379hlzFC6+zlmuuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC1RuAIAAAAAAIAtUbgCAAAAAACALVG4AgAAAAAAgC0FFfQAAAAAUDgcP37cY1a1atV8HAkAIK+9++67XmWS9MYbbxjzffv2eTUmXJq44goAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC25LAsy/K5E4fDH2MB4IEfpqkb5iyQt5izQOHCnAUKF+YsULj4Ome54goAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC2ROEKAAAAAAAAtkThCgAAAAAAALZE4QoAAAAAAAC25LAsyyroQQAAAAAAAAAX4oorAAAAAAAA2BKFKwAAAAAAANgShSsAAAAAAADYEoUrAAAAAAAA2BKFKwAAAAAAANgShSsAAAAAAADYEoUrAAAAAAAA2BKFKwAAAAAAANgShSsAAAAAAADY0v8B5UALYj3Nia8AAAAASUVORK5CYII=", + "text/plain": [ + "
" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "fig, axes = plt.subplots(2, 5, figsize=(12, 5))\n", + "for ax, img, lbl in zip(axes.flat, x_test[:10], y_test[:10]):\n", + " cap = model.caption(jnp.array(img))\n", + " color = 'green' if cap.split()[-1] == DIGIT_NAMES[lbl] else 'red'\n", + " ax.imshow(img, cmap='gray')\n", + " ax.set_title(cap, color=color)\n", + " ax.axis('off')\n", + "plt.tight_layout()" + ] + }, + { + "cell_type": "markdown", + "id": "5d40e8f3", + "metadata": {}, + "source": [ + "The model learns in a few epochs to predict the correct digit word from the image, demonstrating that the patch tokens carry meaningful visual information and that M-RoPE correctly communicates the spatial layout of those patches to the attention layers.\n", + "\n", + "## What to try next\n", + "\n", + "- **Larger images / more patches**: try 14×14 patches on higher-resolution inputs to see M-RoPE's 2D advantage over 1D position embeddings more clearly.\n", + "- **Free-form captions**: replace the templated captions with real image descriptions (e.g. a subset of [COCO Captions](https://cocodataset.org)) and add a pretrained vision encoder (ViT / SigLIP) in place of the raw patch projector.\n", + "- **Model parallelism**: add `PartitionSpec` metadata to the weight initialisers as shown in [miniGPT](minigpt.md) to shard the model across multiple devices." + ] + } + ], + "metadata": { + "jupytext": { + "formats": "md,ipynb", + "main_language": "python" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs_nnx/examples/multimodal_llm.md b/docs_nnx/examples/multimodal_llm.md new file mode 100644 index 000000000..324d7ec7e --- /dev/null +++ b/docs_nnx/examples/multimodal_llm.md @@ -0,0 +1,375 @@ +--- +jupyter: + jupytext: + formats: md,ipynb + main_language: python + text_representation: + extension: .md + format_name: markdown + format_version: '1.3' + jupytext_version: 1.13.8 +--- + +# Build a multimodal language model with JAX + +This tutorial extends [miniGPT](minigpt.md) to handle images alongside text. You will build a small multimodal language model that learns to caption MNIST digits by projecting image patches directly into the token sequence — the same core idea used in [LLaVA](https://arxiv.org/abs/2304.08485). + +The model uses **Multimodal Rotary Position Embeddings (M-RoPE)** from [Qwen2-VL](https://arxiv.org/abs/2409.12191): image tokens are assigned independent 2D spatial positions (row, col) while text tokens get standard 1D sequential positions, with no separate learned positional embedding table. + +Here, you will learn how to: + +- Turn image patches into token embeddings with a linear patch projector +- Assign 2D spatial positions to image tokens and 1D positions to text tokens with M-RoPE +- Build a causal transformer that jointly attends over image and text tokens +- Train the model to generate image captions from scratch + +## Setup + +```python +import jax +import jax.numpy as jnp +import flax.nnx as nnx +import optax +import numpy as np +import matplotlib.pyplot as plt +from tensorboardX import SummaryWriter +import random +import functools +import tiktoken +from datasets import load_dataset +``` + +## Tokenizer + +We use the same GPT-2 tokenizer from [Tiktoken](https://github.com/openai/tiktoken) as in the miniGPT tutorial, giving the model a full-sized vocabulary consistent with a real language model. + +To make training more interesting, each image is paired with a caption drawn at random from several templates. + +```python +tokenizer = tiktoken.get_encoding("gpt2") +BOS = EOS = tokenizer.encode('<|endoftext|>', allowed_special={'<|endoftext|>'})[0] +PAD = 0 +VOCAB_SIZE = tokenizer.n_vocab +``` + +## Dataset: MNIST with captions + +We pair every digit image with a templated caption such as `"the digit is seven"`. The model learns to generate the caption given the image. + +```python +mnist = load_dataset("ylecun/mnist") +x_train = np.stack([np.array(img) for img in mnist['train']['image']]).astype(np.float32) / 255.0 +y_train = np.array(mnist['train']['label']) +x_test = np.stack([np.array(img) for img in mnist['test']['image']]).astype(np.float32) / 255.0 +y_test = np.array(mnist['test']['label']) + +DIGIT_NAMES = ['zero', 'one', 'two', 'three', 'four', + 'five', 'six', 'seven', 'eight', 'nine'] + +CAPTION_TEMPLATES = [ + "the digit is {word}", + "this shows the number {word}", + "I see a {word}", + "the handwritten digit {word}", + "a picture of {word}", + "this digit is {word}", +] + +# TEXT_LEN must cover the longest possible encoded caption (BOS + tokens). +TEXT_LEN = max( + len([BOS] + tokenizer.encode(tmpl.format(word=d))) + for tmpl in CAPTION_TEMPLATES + for d in DIGIT_NAMES +) +``` + +```python +def encode_caption(label: int, template: str | None = None) -> list[int]: + """Encode a digit label using a randomly chosen (or specified) caption template.""" + tmpl = template or random.choice(CAPTION_TEMPLATES) + tokens = [BOS] + tokenizer.encode(tmpl.format(word=DIGIT_NAMES[label])) + return tokens + [PAD] * (TEXT_LEN - len(tokens)) + +def decode_tokens(ids) -> str: + valid = [int(i) for i in ids if int(i) not in (PAD, BOS, EOS)] + return tokenizer.decode(valid) +``` + +The training objective is **next-token prediction** on the text part of the sequence. +Given the image and all previous text tokens, the model predicts the next token: + +| Input | `` | `the` | `digit` | `is` | `seven` | +|--------|---------|---------|---------|---------|---------| +| Target | `the` | `digit` | `is` | `seven` | `` | + +The loss is computed only over the text positions — not over the image patches. + +```python +def make_batch(images: np.ndarray, labels: np.ndarray): + """Build (images, text_inputs, text_targets) JAX arrays for a batch.""" + B = len(labels) + text_in = np.array([encode_caption(lbl) for lbl in labels], dtype=np.int32) + text_out = np.concatenate( + [text_in[:, 1:], np.full((B, 1), EOS, dtype=np.int32)], axis=1 + ) + return jnp.array(images), jnp.array(text_in), jnp.array(text_out) +``` + +## Model architecture + +The model has four components assembled into `MultimodalGPT`: + +1. **`PatchEmbedding`** — splits each 28×28 image into 16 non-overlapping 7×7 patches and projects each flattened patch to `embed_dim` with a single linear layer. +2. **Token embedding** — a standard `nnx.Embed` lookup that maps text token IDs to `embed_dim` vectors. +3. **Causal transformer** — a stack of pre-norm transformer blocks. Each block uses `nnx.MultiHeadAttention` with a custom M-RoPE `attention_fn` so image patches and text tokens share the same sequence but carry different positional coordinates. +4. **LM head** — a linear projection from `embed_dim` to `vocab_size`, applied only to the text positions of the output. The image patch outputs are discarded because the loss is computed purely on text. + +The 16 image tokens are prepended to the text tokens and the combined sequence is processed jointly by the transformer. M-RoPE lets the model distinguish image patches by their 2D spatial location and text tokens by their sequential order. + +### Patch embedding + +We divide each 28×28 image into a 4×4 grid of 7×7 patches, yielding **16 image tokens**. A single linear layer projects each flattened patch to `embed_dim`. + +```python +class PatchEmbedding(nnx.Module): + """Project image patches into the model's embedding space.""" + + def __init__(self, patch_size: int, embed_dim: int, *, rngs: nnx.Rngs): + self.patch_size = patch_size + self.proj = nnx.Linear(patch_size * patch_size, embed_dim, rngs=rngs) + + def __call__(self, images): + # images: (B, H, W) + B, H, W = images.shape + p = self.patch_size + nh, nw = H // p, W // p + x = images.reshape(B, nh, p, nw, p).transpose(0, 1, 3, 2, 4) # (B, nh, nw, p, p) + return self.proj(x.reshape(B, nh * nw, p * p)) # (B, num_patches, embed_dim) +``` + +### Multimodal Rotary Position Embedding (M-RoPE) + +Standard RoPE assigns a single sequential integer to each token. M-RoPE extends this for multimodal sequences by **splitting the attention head dimension into two equal halves**: + +- **First half** is rotated by the *row* index (image patches) or *sequence position* (text). +- **Second half** is rotated by the *column* index (image patches) or *sequence position* (text). + +For text tokens both halves receive the same value, so M-RoPE reduces to standard 1D RoPE. For image patches the two halves carry independent row and column coordinates, encoding 2D location without any additional embedding table. + +``` +Sequence layout: [patch(0,0) patch(0,1) … patch(3,3) | the digit is ] + ←————————————16 image tokens—————————→ ←————5 text tokens————→ + +M-RoPE pos_0: [0 0 0 0 1 1 1 1 2 2 2 2 3 3 3 3 | 4 5 6 7 8] (row / seq) +M-RoPE pos_1: [0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 | 4 5 6 7 8] (col / seq) +``` + +```python +def make_mrope_positions(grid_h: int, grid_w: int, text_len: int): + """Return per-token M-RoPE position IDs as two 1-D arrays. + + Image patches get independent (row, col) positions. Text tokens get a + sequential position in both components, offset past the image grid so + image and text coordinates do not collide. + """ + rows = jnp.repeat(jnp.arange(grid_h), grid_w) + cols = jnp.tile(jnp.arange(grid_w), grid_h) + text_seq = jnp.arange(text_len) + max(grid_h, grid_w) + return jnp.concatenate([rows, text_seq]), jnp.concatenate([cols, text_seq]) +``` + +### Attention, transformer blocks, and the full model + +M-RoPE is implemented with `dot_product_attention_with_mrope`, a thin extension of +`nnx.dot_product_attention_with_rope` for 2D positions. It accepts `input_positions` +as a `[2, seq_len]` array where row 0 holds row/sequence indices and row 1 holds +column/sequence indices. The first and second halves of each head's dimension are +rotated independently by `nnx.RoPE`, then concatenated before the standard +dot-product attention is computed. `nnx.MultiHeadAttention` takes this function as +its `attention_fn`, so all Q/K/V projections and the output projection are handled +by the built-in module. + +```python +def dot_product_attention_with_mrope( + query, key, value, *, rope, input_positions, **kwargs +): + """Dot-product attention with Multimodal RoPE. + + Extends nnx.dot_product_attention_with_rope for 2D spatial positions. + ``input_positions`` must be a [2, seq_len] array: row 0 contains row/sequence + indices (pos_0) and row 1 contains column/sequence indices (pos_1). Image + patches receive independent (row, col) coordinates; text tokens receive the + same value in both rows (standard 1-D RoPE behaviour). + """ + pos_0, pos_1 = input_positions[0], input_positions[1] + h = query.shape[-1] // 2 + apply = jax.vmap(rope, in_axes=(-2, None), out_axes=-2) + query = jnp.concatenate([apply(query[..., :h], pos_0), apply(query[..., h:], pos_1)], axis=-1) + key = jnp.concatenate([apply(key [..., :h], pos_0), apply(key [..., h:], pos_1)], axis=-1) + return nnx.dot_product_attention(query, key, value, **kwargs) +``` + +```python +class CausalSelfAttention(nnx.Module): + """Causal multi-head self-attention with M-RoPE via nnx.MultiHeadAttention.""" + + def __init__(self, embed_dim: int, num_heads: int, *, rngs: nnx.Rngs): + assert embed_dim % num_heads == 0 + assert (embed_dim // num_heads) % 4 == 0, "head_dim must be divisible by 4 for M-RoPE" + self.attn = nnx.MultiHeadAttention( + num_heads=num_heads, + in_features=embed_dim, + use_bias=False, + decode=False, + attention_fn=functools.partial(dot_product_attention_with_mrope, rope=nnx.RoPE()), + rngs=rngs, + ) + + def __call__(self, x, pos_0, pos_1): + return self.attn(x, is_causal=True, input_positions=jnp.stack([pos_0, pos_1])) +``` + +```python +class TransformerBlock(nnx.Module): + """Pre-norm transformer block.""" + + def __init__(self, embed_dim: int, num_heads: int, ff_dim: int, *, rngs: nnx.Rngs): + self.attn = CausalSelfAttention(embed_dim, num_heads, rngs=rngs) + self.fc1 = nnx.Linear(embed_dim, ff_dim, rngs=rngs) + self.fc2 = nnx.Linear(ff_dim, embed_dim, rngs=rngs) + self.norm1 = nnx.LayerNorm(num_features=embed_dim, rngs=rngs) + self.norm2 = nnx.LayerNorm(num_features=embed_dim, rngs=rngs) + + def __call__(self, x, pos_0, pos_1): + x = x + self.attn(self.norm1(x), pos_0, pos_1) + x = x + self.fc2(jax.nn.gelu(self.fc1(self.norm2(x)))) + return x +``` + +```python +class MultimodalGPT(nnx.Module): + """Image-conditioned causal language model. + + Image patches and text tokens are concatenated into a single sequence. + M-RoPE encodes spatial position for patches and sequential position for + text. + """ + + def __init__(self, *, rngs: nnx.Rngs): + self.patch_embed = PatchEmbedding(PATCH_SIZE, EMBED_DIM, rngs=rngs) + self.token_embed = nnx.Embed(num_embeddings=VOCAB_SIZE, features=EMBED_DIM, rngs=rngs) + self.blocks = nnx.List([TransformerBlock(EMBED_DIM, NUM_HEADS, FF_DIM, rngs=rngs) + for _ in range(NUM_BLOCKS)]) + self.norm = nnx.LayerNorm(num_features=EMBED_DIM, rngs=rngs) + self.lm_head = nnx.Linear(EMBED_DIM, VOCAB_SIZE, use_bias=False, rngs=rngs) + + def __call__(self, images, text_tokens): + """ + Args: + images: (B, H, W) float32 in [0, 1] + text_tokens: (B, text_len) int32 + Returns: + logits: (B, text_len, vocab_size) + """ + text_len = text_tokens.shape[1] + x = jnp.concatenate([ + self.patch_embed(images), # (B, num_patches, embed_dim) + self.token_embed(text_tokens), # (B, text_len, embed_dim) + ], axis=1) + + pos_0, pos_1 = make_mrope_positions(GRID_H, GRID_W, text_len) + for block in self.blocks: + x = block(x, pos_0, pos_1) + + return self.lm_head(self.norm(x)[:, GRID_H * GRID_W:]) # text positions only + + def caption(self, image): + """Greedily decode a caption for a single (H, W) image.""" + tokens = [BOS] + image = image[None] # add batch dim + for _ in range(TEXT_LEN - 1): + text_in = jnp.array([tokens + [PAD] * (TEXT_LEN - len(tokens))]) + next_id = int(jnp.argmax(self(image, text_in)[0, len(tokens) - 1])) + if next_id == EOS: + break + tokens.append(next_id) + return decode_tokens(tokens[1:]) # strip BOS +``` + +## Set hyperparameters and create the model + +```python +PATCH_SIZE = 7 # 28 / 7 = 4 patches per side +GRID_H = GRID_W = 28 // PATCH_SIZE +EMBED_DIM = 64 +NUM_HEADS = 4 # head_dim = 16, divisible by 4 ✓ +FF_DIM = 256 +NUM_BLOCKS = 4 +BATCH_SIZE = 256 +NUM_EPOCHS = 5 +LR = 1e-3 +``` + +```python +model = MultimodalGPT(rngs=nnx.Rngs(0)) +optimizer = nnx.Optimizer(model, optax.adam(LR), wrt=nnx.Param) +``` + +## Define the loss function and training step + +The training objective is **next-token prediction** on the text part of the sequence. `loss_fn` runs a forward pass and computes the mean cross-entropy between the model's logits (shape `[B, text_len, vocab_size]`) and the shifted target token IDs from `make_batch`. Image patch positions are excluded from the loss — the model only needs to learn to predict text given the visual context. + +`train_step` wraps `loss_fn` with `nnx.value_and_grad` to obtain both the scalar loss and the parameter gradients in one call, then applies the update through the Optax optimizer. Decorating the step with `@nnx.jit` compiles the entire forward-backward-update pass with XLA so it runs efficiently on GPU and TPU. + +```python +def loss_fn(model, images, text_in, text_out): + logits = model(images, text_in) + return optax.softmax_cross_entropy_with_integer_labels(logits, text_out).mean() + +@nnx.jit +def train_step(model, optimizer, images, text_in, text_out): + loss, grads = nnx.value_and_grad(loss_fn)(model, images, text_in, text_out) + optimizer.update(model, grads) + return loss +``` + +## Train the model + +```python +writer = SummaryWriter() +step = 0 +for epoch in range(NUM_EPOCHS): + perm = np.random.permutation(len(x_train)) + epoch_losses = [] + + for i in range(0, len(x_train) - BATCH_SIZE + 1, BATCH_SIZE): + idx = perm[i : i + BATCH_SIZE] + images, text_in, text_out = make_batch(x_train[idx], y_train[idx]) + loss = float(train_step(model, optimizer, images, text_in, text_out)) + writer.add_scalar('train_loss', loss, step) + epoch_losses.append(loss) + step += 1 + + print(f"Epoch {epoch + 1}/{NUM_EPOCHS} loss={np.mean(epoch_losses):.4f}") +``` + +## Evaluate + +```python +fig, axes = plt.subplots(2, 5, figsize=(12, 5)) +for ax, img, lbl in zip(axes.flat, x_test[:10], y_test[:10]): + cap = model.caption(jnp.array(img)) + color = 'green' if cap.split()[-1] == DIGIT_NAMES[lbl] else 'red' + ax.imshow(img, cmap='gray') + ax.set_title(cap, color=color) + ax.axis('off') +plt.tight_layout() +``` + +The model learns in a few epochs to predict the correct digit word from the image, demonstrating that the patch tokens carry meaningful visual information and that M-RoPE correctly communicates the spatial layout of those patches to the attention layers. + +## What to try next + +- **Larger images / more patches**: try 14×14 patches on higher-resolution inputs to see M-RoPE's 2D advantage over 1D position embeddings more clearly. +- **Free-form captions**: replace the templated captions with real image descriptions (e.g. a subset of [COCO Captions](https://cocodataset.org)) and add a pretrained vision encoder (ViT / SigLIP) in place of the raw patch projector. +- **Model parallelism**: add `PartitionSpec` metadata to the weight initialisers as shown in [miniGPT](minigpt.md) to shard the model across multiple devices.