code: def _speculative_sampling(
      candidate_input_ids,
      candidate_logits,
      candidate_length,
      new_logits,
      last_assistant_token_is_eos,
      max_matches,
  ):
      """
      Applies sampling as in the speculative decoding paper (https://arxiv.org/pdf/2211.17192.pdf, algorithm 1). Returns
      the selected tokens, as well as the number of candidate matches.
  
      NOTE: Unless otherwise stated, the variable names match those in the paper.
      """
      # Gets the probabilities from the logits. q_i and p_i denote the assistant and model probabilities of the tokens
      # selected by the assistant, respectively.
      q = candidate_logits.softmax(dim=-1)
      q_i = q[
          :,
          torch.arange(candidate_length),
          candidate_input_ids[:, -candidate_length:],
      ].squeeze(0, 1)
      p = new_logits.softmax(dim=-1)
      p_i = p[
          :,
          torch.arange(candidate_length),
          candidate_input_ids[:, -candidate_length:],
      ].squeeze(0, 1)
      probability_ratio = p_i / q_i
  
      # When probability_ratio > 1 (i.e. q_i(x) < p_i(x), or "assistant probability of the candidate token is smaller
      # than the model probability for the same token"), keep the token. Otherwise reject with p = 1 - probability_ratio
      # (= keep with p = probability_ratio). Keep all the tokens until the first rejection
      r_i = torch.rand_like(probability_ratio)
      is_accepted = r_i <= probability_ratio
      n_matches = (
          ~is_accepted.cumsum(dim=-1) < 1
      ).sum()  # this is `n` in algorithm 1
  
      # Ensure we don't generate beyond max_len or an EOS token (not in algorithm 1, but needed for correct behavior)
      if last_assistant_token_is_eos and n_matches == candidate_length:
          n_matches -= 1
      n_matches = min(n_matches, max_matches)
  
      # Next token selection: if there is a rejection, adjust the distribution from the main model before sampling.
      gamma = candidate_logits.shape[1]
      p_n_plus_1 = p[:, n_matches, :]
      if n_matches < gamma:
          q_n_plus_1 = q[:, n_matches, :]
          p_prime = torch.clamp(
              (p_n_plus_1 - q_n_plus_1), min=0
          ).softmax(dim=-1)
      else:
          p_prime = p_n_plus_1
      t = torch.multinomial(p_prime, num_samples=1).squeeze(1)[None, :]
  
      # The selected tokens include the matches (if any) plus the next sampled tokens
      if n_matches > 0:
          valid_tokens = torch.cat(
              (candidate_input_ids[:, -n_matches:], t), dim=-1
          )
      else:
          valid_tokens = t
  
      return valid_tokens, n_matches
  
imports: import copy
  import inspect
  import warnings
  from dataclasses import dataclass
  from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
  import torch
  import torch.distributed
  from torch import nn
  from ..cache_utils import Cache, DynamicCache
  from ..integrations.deepspeed import is_deepspeed_zero3_enabled
  from ..modeling_outputs import CausalLMOutputWithPast, Seq2SeqLMOutput
  from ..models.auto import MODEL_FOR_CAUSAL_IMAGE_MODELING_MAPPING, MODEL_FOR_CAUSAL_LM_MAPPING, MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING, MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING, MODEL_FOR_VISION_2_SEQ_MAPPING
  from ..utils import ExplicitEnum, ModelOutput, is_accelerate_available, logging
  from .beam_constraints import DisjunctiveConstraint, PhrasalConstraint
  from .beam_search import BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer
  from .candidate_generator import AssistedCandidateGenerator, CandidateGenerator, PromptLookupCandidateGenerator, _crop_past_key_values, _prepare_attention_mask, _prepare_token_type_ids
  from .configuration_utils import GenerationConfig
  from .logits_process import EncoderNoRepeatNGramLogitsProcessor, EncoderRepetitionPenaltyLogitsProcessor, EpsilonLogitsWarper, EtaLogitsWarper, ExponentialDecayLengthPenalty, ForcedBOSTokenLogitsProcessor, ForcedEOSTokenLogitsProcessor, ForceTokensLogitsProcessor, HammingDiversityLogitsProcessor, InfNanRemoveLogitsProcessor, LogitNormalization, LogitsProcessorList, MinLengthLogitsProcessor, MinNewTokensLengthLogitsProcessor, NoBadWordsLogitsProcessor, NoRepeatNGramLogitsProcessor, PrefixConstrainedLogitsProcessor, RepetitionPenaltyLogitsProcessor, SequenceBiasLogitsProcessor, SuppressTokensAtBeginLogitsProcessor, SuppressTokensLogitsProcessor, TemperatureLogitsWarper, TopKLogitsWarper, TopPLogitsWarper, TypicalLogitsWarper, UnbatchedClassifierFreeGuidanceLogitsProcessor
  from .stopping_criteria import MaxLengthCriteria, MaxTimeCriteria, StoppingCriteria, StoppingCriteriaList, validate_stopping_criteria
label_desc: Use image node_id 1 for calling a global function with example usage: _speculative_sampling(candidate_input_ids, candidate_logits, candidate_length, new_logits, last_assistant_token_is_eos, max_matches) and returns: valid_tokens,  n_matches