Merge remote-tracking branch 'upstream/master' into jax-native

anthony-etim · anthony-etim · commit cd7617f09743 · 2026-07-24T07:36:41.000Z
diff --git a/.github/workflows/actions.yml b/.github/workflows/actions.yml
@@ -29,28 +29,25 @@ jobs:
             version: keras-nightly
           - backend: openvino
             version: keras-nightly
-    runs-on: ubuntu-latest
+
+    container:
+      image: python:3.11-slim
+
+    runs-on: linux-x86-n2-16
+    timeout-minutes: 120
     env:
       KERAS_BACKEND: ${{ matrix.backend }}
     steps:
-    - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
-      with:
-        persist-credentials: false
-    - name: Set up Python 3.11
-      uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
-      with:
-        python-version: '3.11'
-    - name: Get pip cache dir
-      id: pip-cache
+    - name: Install binary dependencies
+      # build-essential for C extensions (tokenizers, torch)
+      # curl, git, gnupg for checkout and other actions
       run: |
-        python -m pip install --upgrade pip setuptools
-        echo "::set-output name=dir::$(pip cache dir)"
-    - name: pip cache
-      # zizmor: ignore[cache-poisoning]
-      uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
+        apt-get update
+        apt-get -y install build-essential curl git gnupg
+        git config --global --add safe.directory '*'
+    - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
       with:
-        path: ${{ steps.pip-cache.outputs.dir }}
-        key: ${{ runner.os }}-pip-${{ hashFiles('requirements-common.txt') }}-${{ hashFiles('requirements.txt') }}-${{ hashFiles('pyproject.toml') }}
+        persist-credentials: false
     - name: Install dependencies
       run: |
           pip install -r requirements.txt --progress-bar off
@@ -68,7 +65,7 @@ jobs:
         pip install keras-nightly --no-cache-dir --progress-bar off
     - name: Test with pytest
       run: |
-        pytest keras_hub/
+        pytest keras_hub/ -n auto --dist loadfile --durations=50
     - name: Run integration tests
       run: |
         python pip_build.py --install
@@ -82,7 +79,7 @@ jobs:
     name: Check the code format
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
+      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
         with:
           persist-credentials: false
       - name: Set up Python 3.11
diff --git a/.github/workflows/auto-assignment.yml b/.github/workflows/auto-assignment.yml
@@ -16,11 +16,11 @@ jobs:
   welcome:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
+      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
         with:
           persist-credentials: false
           ref: ${{ github.event.repository.default_branch }}
-      - uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9
+      - uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
         with:
           script: |
             const script = require('./\.github/workflows/scripts/auto-assignment.js')
diff --git a/.github/workflows/gemini-automated-issue-triage.yml b/.github/workflows/gemini-automated-issue-triage.yml
@@ -55,7 +55,7 @@ jobs:
         id: 'get_issue_data'
         if: |-
           github.event_name == 'workflow_dispatch'
-        uses: 'actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3' # v9
+        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
         env:
           ISSUE_NUMBER: '${{ github.event.inputs.issue_number || inputs.issue_number }}'
         with:
@@ -73,13 +73,13 @@ jobs:
             return issue;
 
       - name: 'Checkout'
-        uses: 'actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0' # v7
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
         with:
           persist-credentials: false
 
       - name: 'Get Repository Labels'
         id: 'get_labels'
-        uses: 'actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3' # v9
+        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
         with:
           github-token: '${{ secrets.GITHUB_TOKEN }}'
           script: |-
@@ -115,7 +115,7 @@ jobs:
             ${{ github.event_name == 'workflow_dispatch' && steps.get_issue_data.outputs.title || github.event.issue.title }}
           ISSUE_BODY: >-
             ${{ github.event_name == 'workflow_dispatch' && steps.get_issue_data.outputs.body || github.event.issue.body }}
-        uses: 'actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3' # v9
+        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
         with:
           github-token: '${{ secrets.GITHUB_TOKEN }}'
           script: |-
@@ -160,7 +160,7 @@ jobs:
         run: mkdir -p ~/.gemini
 
       - name: 'Run Gemini Issue Analysis & Response'
-        uses: 'google-github-actions/run-gemini-cli@f77273f4c914e4bf38440cf36a0369cb64a37489'
+        uses: google-github-actions/run-gemini-cli@f77273f4c914e4bf38440cf36a0369cb64a37489
         id: 'gemini_issue_analysis'
         env:
           GITHUB_TOKEN: '' # Do not pass any auth token here
@@ -233,7 +233,7 @@ jobs:
           ISSUE_NUMBER: '${{ github.event.issue.number || github.event.inputs.issue_number }}'
           LABELS_OUTPUT: '${{ steps.gemini_issue_analysis.outputs.summary }}'
           ALLOWED_LABELS: '${{ steps.get_labels.outputs.available_labels }}'
-        uses: 'actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3' # v9
+        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
         with:
           github-token: '${{ secrets.GITHUB_TOKEN }}'
           script: |
diff --git a/.github/workflows/labeler.yaml b/.github/workflows/labeler.yaml
@@ -14,11 +14,11 @@ jobs:
   welcome:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
+      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
         with:
           persist-credentials: false
           ref: ${{ github.event.repository.default_branch }}
-      - uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9
+      - uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
         with:
           script: |
             const script = require('./\.github/workflows/scripts/labeler.js')
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
@@ -16,7 +16,7 @@ jobs:
     needs: [run-test-for-nightly]
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
+      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
         with:
           persist-credentials: false
       - name: Set up Python
diff --git a/.github/workflows/publish-to-pypi.yml b/.github/workflows/publish-to-pypi.yml
@@ -10,7 +10,7 @@ jobs:
     name: Build and publish to PyPI
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
+      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
         with:
           persist-credentials: false
       - name: Set up Python
diff --git a/.github/workflows/scorecard.yml b/.github/workflows/scorecard.yml
@@ -25,7 +25,7 @@ jobs:
 
     steps:
       - name: "Checkout code"
-        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
         with:
           persist-credentials: false
 
diff --git a/.github/workflows/zizmor.yml b/.github/workflows/zizmor.yml
@@ -14,7 +14,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout code
-        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
         with:
           persist-credentials: false
 
diff --git a/keras_hub/api/models/__init__.py b/keras_hub/api/models/__init__.py
@@ -511,6 +511,9 @@
 from keras_hub.src.models.mistral.mistral_causal_lm_preprocessor import (
     MistralCausalLMPreprocessor as MistralCausalLMPreprocessor,
 )
+from keras_hub.src.models.mistral.mistral_tokenizer import (
+    MistralTekkenTokenizer as MistralTekkenTokenizer,
+)
 from keras_hub.src.models.mistral.mistral_tokenizer import (
     MistralTokenizer as MistralTokenizer,
 )
diff --git a/keras_hub/api/tokenizers/__init__.py b/keras_hub/api/tokenizers/__init__.py
@@ -71,6 +71,9 @@
 from keras_hub.src.models.metaclip_2.metaclip_2_tokenizer import (
     MetaCLIP2Tokenizer as MetaCLIP2Tokenizer,
 )
+from keras_hub.src.models.mistral.mistral_tokenizer import (
+    MistralTekkenTokenizer as MistralTekkenTokenizer,
+)
 from keras_hub.src.models.mistral.mistral_tokenizer import (
     MistralTokenizer as MistralTokenizer,
 )
diff --git a/keras_hub/src/models/mistral/mistral_attention.py b/keras_hub/src/models/mistral/mistral_attention.py
@@ -27,13 +27,15 @@ def __init__(
         kernel_initializer="glorot_uniform",
         sliding_window=512,
         dropout=0,
+        head_dim=None,
         **kwargs,
     ):
         super().__init__(**kwargs)
         self._num_query_heads = num_query_heads
         self._num_key_value_heads = num_key_value_heads
         self._sliding_window = sliding_window
         self._dropout = dropout
+        self._head_dim = head_dim
 
         self._num_key_value_groups = num_query_heads // num_key_value_heads
         self._rope_max_wavelength = rope_max_wavelength
@@ -54,7 +56,8 @@ def build(self, inputs_shape):
         # v = num key/value heads
         # h = head dim
         self._hidden_dim = inputs_shape[-1]
-        self._head_dim = self._hidden_dim // self._num_query_heads
+        if self._head_dim is None:
+            self._head_dim = self._hidden_dim // self._num_query_heads
         self._inv_norm_factor = 1.0 / math.sqrt(self._head_dim)
 
         self._query_dense = keras.layers.EinsumDense(
@@ -239,6 +242,7 @@ def get_config(self):
                 ),
                 "sliding_window": self._sliding_window,
                 "dropout": self._dropout,
+                "head_dim": self._head_dim,
             }
         )
         return config
diff --git a/keras_hub/src/models/mistral/mistral_backbone.py b/keras_hub/src/models/mistral/mistral_backbone.py
@@ -52,7 +52,13 @@ class MistralBackbone(Backbone):
             attention layers. This controls the maximum cache size for the
             attention layers in each transformer decoder. Only `sliding_window`
             number of tokens are saved in the cache and used to generate the
-            next token. Defaults to `512`.
+            next token. Defaults to `512`. Pass `None` to disable sliding
+            window attention entirely (e.g. Magistral).
+        head_dim (int, optional): The size of each attention head. When
+            `None` (the default), falls back to `hidden_dim // num_query_heads`.
+            Set explicitly when the model's head size is not equal to
+            `hidden_dim // num_query_heads` — e.g. Magistral uses
+            `head_dim=128` with `hidden_dim=5120` and `num_query_heads=32`.
         dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
             for model computations and weights. Note that some computations,
             such as softmax and layer normalization, will always be done at
@@ -98,6 +104,7 @@ def __init__(
         rope_scaling_factor=1.0,
         layer_norm_epsilon=1e-6,
         sliding_window=512,
+        head_dim=None,
         dropout=0,
         dtype=None,
         **kwargs,
@@ -123,6 +130,7 @@ def __init__(
                 activation=ops.silu,
                 kernel_initializer=_mistral_kernel_initializer(stddev=0.02),
                 sliding_window=sliding_window,
+                head_dim=head_dim,
                 dropout=dropout,
                 dtype=dtype,
                 name=f"transformer_layer_{i}",
@@ -165,6 +173,7 @@ def __init__(
         self.num_key_value_heads = num_key_value_heads
         self.rope_scaling_factor = rope_scaling_factor
         self.sliding_window = sliding_window
+        self.head_dim = head_dim
         self.layer_norm_epsilon = layer_norm_epsilon
         self.dropout = dropout
 
@@ -181,6 +190,7 @@ def get_config(self):
                 "rope_scaling_factor": self.rope_scaling_factor,
                 "num_key_value_heads": self.num_key_value_heads,
                 "sliding_window": self.sliding_window,
+                "head_dim": self.head_dim,
                 "layer_norm_epsilon": self.layer_norm_epsilon,
                 "dropout": self.dropout,
             }
diff --git a/keras_hub/src/models/mistral/mistral_backbone_test.py b/keras_hub/src/models/mistral/mistral_backbone_test.py
@@ -42,6 +42,26 @@ def test_num_parameters(self):
         # Reference value calculated using the PyTorch model
         self.assertEqual(model.count_params(), 2704)
 
+    def test_explicit_head_dim(self):
+        # Magistral-style config: `head_dim` is set explicitly and does not
+        # equal `hidden_dim // num_query_heads`. `sliding_window=None` is
+        # also exercised here. Run the full backbone test so the new path
+        # gets serialization and precision coverage.
+        init_kwargs = {
+            **self.init_kwargs,
+            "sliding_window": None,
+            "head_dim": 4,
+        }
+        self.run_backbone_test(
+            cls=MistralBackbone,
+            init_kwargs=init_kwargs,
+            input_data=self.input_data,
+            expected_output_shape=(2, 5, 16),
+        )
+        model = MistralBackbone(**init_kwargs)
+        attention = model.transformer_layers[0]._self_attention_layer
+        self.assertEqual(attention._head_dim, 4)
+
     @pytest.mark.extra_large
     def test_smallest_preset(self):
         self.run_preset_test(
diff --git a/keras_hub/src/models/mistral/mistral_causal_lm.py b/keras_hub/src/models/mistral/mistral_causal_lm.py
@@ -100,7 +100,9 @@ def _build_cache(self, token_ids):
         max_length = ops.shape(token_ids)[1]
         num_layers = self.backbone.num_layers
         num_key_value_heads = self.backbone.num_key_value_heads
-        head_dim = self.backbone.hidden_dim // self.backbone.num_query_heads
+        head_dim = self.backbone.head_dim or (
+            self.backbone.hidden_dim // self.backbone.num_query_heads
+        )
         shape = [
             batch_size,
             num_layers,
diff --git a/keras_hub/src/models/mistral/mistral_causal_lm_test.py b/keras_hub/src/models/mistral/mistral_causal_lm_test.py
@@ -47,6 +47,30 @@ def test_causal_lm_basics(self):
             expected_output_shape=(2, 8, 10),
         )
 
+    def test_generate_with_explicit_head_dim(self):
+        # Magistral-style config where `head_dim` is set explicitly and does
+        # not equal `hidden_dim // num_query_heads`. The generation cache must
+        # be built with the explicit `head_dim`, not the derived value.
+        backbone = MistralBackbone(
+            vocabulary_size=self.preprocessor.tokenizer.vocabulary_size(),
+            num_layers=2,
+            num_query_heads=4,
+            num_key_value_heads=2,
+            hidden_dim=8,
+            intermediate_dim=16,
+            head_dim=6,
+            sliding_window=None,
+        )
+        self.assertNotEqual(
+            backbone.head_dim, backbone.hidden_dim // backbone.num_query_heads
+        )
+        causal_lm = MistralCausalLM(
+            preprocessor=self.preprocessor, backbone=backbone
+        )
+        prompt = "the quick brown fox"
+        output = causal_lm.generate(prompt)
+        self.assertTrue(prompt in output)
+
     def test_generate(self):
         causal_lm = MistralCausalLM(**self.init_kwargs)
         # String input.
diff --git a/keras_hub/src/models/mistral/mistral_presets.py b/keras_hub/src/models/mistral/mistral_presets.py
@@ -2,6 +2,7 @@
 
 # Metadata for loading pretrained model weights.
 backbone_presets = {
+    # Mistral models
     "mistral_7b_en": {
         "metadata": {
             "description": "Mistral 7B base model",
@@ -42,4 +43,21 @@
         },
         "kaggle_handle": "kaggle://keras/mistral/keras/mistral_0.3_instruct_7b_en/1",
     },
+    # Magistral models
+    "magistral_small_2506_en": {
+        "metadata": {
+            "description": "Magistral Small 2506 model",
+            "params": 23572403200,
+            "path": "mistral",
+        },
+        "kaggle_handle": "kaggle://keras/magistral/keras/magistral_small_2506_en/1",
+    },
+    "magistral_small_2507_en": {
+        "metadata": {
+            "description": "Magistral Small 2507 model",
+            "params": 23572403200,
+            "path": "mistral",
+        },
+        "kaggle_handle": "kaggle://keras/magistral/keras/magistral_small_2507_en/1",
+    },
 }
diff --git a/keras_hub/src/models/mistral/mistral_tokenizer.py b/keras_hub/src/models/mistral/mistral_tokenizer.py
diff --git a/keras_hub/src/models/mistral/mistral_tokenizer_test.py b/keras_hub/src/models/mistral/mistral_tokenizer_test.py
diff --git a/keras_hub/src/models/mistral/mistral_transformer_decoder.py b/keras_hub/src/models/mistral/mistral_transformer_decoder.py
diff --git a/keras_hub/src/utils/transformers/convert_mistral.py b/keras_hub/src/utils/transformers/convert_mistral.py
diff --git a/keras_hub/src/utils/transformers/convert_mistral_test.py b/keras_hub/src/utils/transformers/convert_mistral_test.py
diff --git a/requirements-common.txt b/requirements-common.txt
diff --git a/tools/checkpoint_conversion/convert_mistral_checkpoints.py b/tools/checkpoint_conversion/convert_mistral_checkpoints.py

Original file line number	Diff line number	Diff line change
`@@ -511,6 +511,9 @@`
`511`	`511`	`from keras_hub.src.models.mistral.mistral_causal_lm_preprocessor import (`
`512`	`512`	`MistralCausalLMPreprocessor as MistralCausalLMPreprocessor,`
`513`	`513`	`)`
	`514`	`+from keras_hub.src.models.mistral.mistral_tokenizer import (`
	`515`	`+ MistralTekkenTokenizer as MistralTekkenTokenizer,`
	`516`	`+)`
`514`	`517`	`from keras_hub.src.models.mistral.mistral_tokenizer import (`
`515`	`518`	`MistralTokenizer as MistralTokenizer,`
`516`	`519`	`)`
Original file line number	Diff line number	Diff line change
`@@ -71,6 +71,9 @@`
`71`	`71`	`from keras_hub.src.models.metaclip_2.metaclip_2_tokenizer import (`
`72`	`72`	`MetaCLIP2Tokenizer as MetaCLIP2Tokenizer,`
`73`	`73`	`)`
	`74`	`+from keras_hub.src.models.mistral.mistral_tokenizer import (`
	`75`	`+ MistralTekkenTokenizer as MistralTekkenTokenizer,`
	`76`	`+)`
`74`	`77`	`from keras_hub.src.models.mistral.mistral_tokenizer import (`
`75`	`78`	`MistralTokenizer as MistralTokenizer,`
`76`	`79`	`)`