Repository: huggingface/transformers
Pull request #31755 (2024-07-02 17:16:43)

Requires for torch.tensor before casting
Fixes ONNX export for swin, swin-donut and clap models

https://github.com/huggingface/transformers/blob/82486e5995ed0a65520b10ce1ea938214a199231/src/transformers/models/swin/modeling_swin.py#L643

coming from : 

```
self.shift_size = torch_int(0)
```

as `torch_int` is expecting a `torch.Tensor` : https://github.com/huggingface/transformers/blob/82486e5995ed0a65520b10ce1ea938214a199231/src/transformers/utils/generic.py#L765


also I think we should be able to have here 
```
self.shift_size = 0
```

cc @merveenoyan @xenova 





HuggingFaceDocBuilderDev: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_31755). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Patch:
diff --git a/docs/source/en/model_doc/gemma2.md b/docs/source/en/model_doc/gemma2.md
index fa16dfbc4b..5befa0b1f4 100644
--- a/docs/source/en/model_doc/gemma2.md
+++ b/docs/source/en/model_doc/gemma2.md
@@ -19,12 +19,12 @@ rendered properly in your Markdown viewer.
 
 ## Overview
 
-The Gemma2 model was proposed in [Gemma2: Open Models Based on Gemini Technology and Research](https://blog.google/technology/developers/Gemma2-open-models/) by Gemma2 Team, Google.
-Gemma2 models are trained on 6T tokens, and released with 2 versions, 2b and 7b.
+The Gemma2 model was proposed in [Gemma2: Open Models Based on Gemini Technology and Research](https://blog.google/technology/developers/google-gemma-2/) by Gemma2 Team, Google.
+Two Gemma2 models are released, with parameters sizes of 9 billion (9B) and 27 billion (27B).
 
-The abstract from the paper is the following:
+The abstract from the blog post is the following:
 
-*This work introduces Gemma2, a new family of open language models demonstrating strong performance across academic benchmarks for language understanding, reasoning, and safety. We release two sizes of models (2 billion and 7 billion parameters), and provide both pretrained and fine-tuned checkpoints. Gemma2 outperforms similarly sized open models on 11 out of 18 text-based tasks, and we present comprehensive evaluations of safety and responsibility aspects of the models, alongside a detailed description of our model development. We believe the responsible release of LLMs is critical for improving the safety of frontier models, and for enabling the next wave of LLM innovations*
+*Now we’re officially releasing Gemma 2 to researchers and developers globally. Available in both 9 billion (9B) and 27 billion (27B) parameter sizes, Gemma 2 is higher-performing and more efficient at inference than the first generation, with significant safety advancements built in. In fact, at 27B, it offers competitive alternatives to models more than twice its size, delivering the kind of performance that was only possible with proprietary models as recently as December.*
 
 Tips:
 
diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py
index 25ec7be1b5..f99ae64fb8 100644
--- a/src/transformers/generation/utils.py
+++ b/src/transformers/generation/utils.py
@@ -1493,8 +1493,11 @@ class GenerationMixin:
                 device = self.device
 
             token = token_kwargs if token_kwargs is not None else token_self
-            if token is None or isinstance(token, torch.Tensor):
+            if token is None:
                 return token
+            elif isinstance(token, torch.Tensor):
+                return token.to(device)
+
             return torch.tensor(token, device=device, dtype=torch.long)
 
         bos_token_id = _tensor_or_none(
diff --git a/src/transformers/utils/chat_template_utils.py b/src/transformers/utils/chat_template_utils.py
index ee6173f2a1..414d2fb724 100644
--- a/src/transformers/utils/chat_template_utils.py
+++ b/src/transformers/utils/chat_template_utils.py
@@ -80,7 +80,7 @@ def _parse_type_hint(hint: str) -> Dict:
             return_dict = subtypes[0]
         elif all(isinstance(subtype["type"], str) for subtype in subtypes):
             # A union of basic types can be expressed as a list in the schema
-            return_dict = {"type": [subtype["type"] for subtype in subtypes]}
+            return_dict = {"type": sorted([subtype["type"] for subtype in subtypes])}
         else:
             # A union of more complex types requires "anyOf"
             return_dict = {"anyOf": subtypes}
diff --git a/src/transformers/utils/generic.py b/src/transformers/utils/generic.py
index 80232898ce..01c5ede34a 100644
--- a/src/transformers/utils/generic.py
+++ b/src/transformers/utils/generic.py
@@ -762,7 +762,7 @@ def torch_int(x):
 
     import torch
 
-    return x.to(torch.int64) if torch.jit.is_tracing() else int(x)
+    return x.to(torch.int64) if torch.jit.is_tracing() and isinstance(x, torch.Tensor) else int(x)
 
 
 def torch_float(x):
@@ -774,7 +774,7 @@ def torch_float(x):
 
     import torch
 
-    return x.to(torch.float32) if torch.jit.is_tracing() else int(x)
+    return x.to(torch.float32) if torch.jit.is_tracing() and isinstance(x, torch.Tensor) else int(x)
 
 
 def filter_out_non_signature_kwargs(extra: Optional[list] = None):


Test patch:
diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py
index b9e962a6a1..8fa41fbdbe 100644
--- a/tests/generation/test_utils.py
+++ b/tests/generation/test_utils.py
@@ -30,7 +30,9 @@ from transformers.testing_utils import (
     require_auto_gptq,
     require_quanto,
     require_torch,
+    require_torch_gpu,
     require_torch_multi_accelerator,
+    require_torch_multi_gpu,
     slow,
     torch_device,
 )
@@ -3097,6 +3099,54 @@ class GenerationIntegrationTests(unittest.TestCase, GenerationIntegrationTestsMi
         self.assertTrue(y_prob > 0.001 and n_prob > 0.001)
         self.assertTrue(y_prob <= 1.0 and n_prob <= 1.0)
 
+    @slow
+    @require_torch_multi_gpu
+    def test_assisted_decoding_in_different_gpu(self):
+        # PT-only test: TF doesn't support assisted decoding yet.
+        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM").to("cuda:0")
+        assistant = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM").to(
+            "cuda:1"
+        )
+        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM")
+        model.config.pad_token_id = tokenizer.eos_token_id
+        assistant.config.pad_token_id = tokenizer.eos_token_id
+
+        text = "Hello world"
+        tokenized_inputs = tokenizer([text], return_tensors="pt")
+        input_ids = tokenized_inputs.input_ids.to(torch_device)
+        input_length = input_ids.shape[-1]
+
+        out = model.generate(
+            input_ids,
+            assistant_model=assistant,
+            max_new_tokens=20,
+        )
+        self.assertTrue(input_length <= out.shape[-1] <= input_length + 20)
+
+    @slow
+    @require_torch_gpu
+    def test_assisted_decoding_in_gpu_cpu(self):
+        # PT-only test: TF doesn't support assisted decoding yet.
+        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM").to("cuda")
+        assistant = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM").to(
+            "cpu"
+        )
+        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM")
+        model.config.pad_token_id = tokenizer.eos_token_id
+        assistant.config.pad_token_id = tokenizer.eos_token_id
+
+        text = "Hello world"
+        tokenized_inputs = tokenizer([text], return_tensors="pt")
+        input_ids = tokenized_inputs.input_ids.to(torch_device)
+        input_length = input_ids.shape[-1]
+
+        out = model.generate(
+            input_ids,
+            assistant_model=assistant,
+            max_new_tokens=20,
+        )
+        self.assertTrue(input_length <= out.shape[-1] <= input_length + 20)
+
 
 @require_torch
 class TokenHealingTestCase(unittest.TestCase):
diff --git a/tests/utils/test_chat_template_utils.py b/tests/utils/test_chat_template_utils.py
index cff31c1f8a..1816ddd951 100644
--- a/tests/utils/test_chat_template_utils.py
+++ b/tests/utils/test_chat_template_utils.py
@@ -137,7 +137,7 @@ class JsonSchemaGeneratorTest(unittest.TestCase):
                 "properties": {
                     "x": {
                         "type": "array",
-                        "items": {"type": "array", "items": {"type": ["string", "integer"]}},
+                        "items": {"type": "array", "items": {"type": ["integer", "string"]}},
                         "description": "The input",
                     }
                 },
@@ -455,13 +455,13 @@ class JsonSchemaGeneratorTest(unittest.TestCase):
                     },
                     "y": {
                         "type": "array",
-                        "items": {"type": ["string", "integer"]},
+                        "items": {"type": ["integer", "string"]},
                         "nullable": True,
                         "description": "The second input. It's a big list with a single-line description.",
                     },
                     "z": {
                         "type": "array",
-                        "prefixItems": [{"type": ["string", "integer"]}, {"type": "string"}],
+                        "prefixItems": [{"type": ["integer", "string"]}, {"type": "string"}],
                         "description": "The third input. It's some kind of tuple with a default arg.",
                     },
                 },