code: def compute_contrastive_loss(
      quantized_features,
      transformer_features,
      negative_indices,
      mask_time_indices,
      logits_temp,
      num_negatives,
  ):
      (
          batch_size,
          sequence_length,
          hidden_size,
      ) = quantized_features.shape
  
      # take negative vectors from sampled indices
      quantized_negatives = quantized_features.reshape(-1, hidden_size)[
          negative_indices.reshape(-1)
      ]
      quantized_negatives = quantized_negatives.reshape(
          batch_size, sequence_length, num_negatives, hidden_size
      ).transpose(2, 0, 1, 3)
  
      target_features = jnp.concatenate(
          [quantized_features[None, :], quantized_negatives], axis=0
      )
      loss_logits = optax.cosine_similarity(
          transformer_features, target_features
      )
      loss_logits = loss_logits / logits_temp
  
      neg_is_pos = (quantized_features == quantized_negatives).all(-1)
      neg_is_pos = jnp.concatenate(
          [jnp.full((1,) + loss_logits.shape[1:], False), neg_is_pos],
          axis=0,
      )
  
      # make sure incorrectly sampled vectors don't contribute to loss
      loss_logits = jnp.where(neg_is_pos, -1e9, loss_logits)
  
      predictions = loss_logits.transpose(2, 1, 0).reshape(
          -1, loss_logits.shape[0]
      )
      targets = (
          ((1 - mask_time_indices) * -100).transpose(1, 0).flatten()
      )
  
      target_mask = jnp.where(targets >= 0, 1.0, 0.0)
      contrastive_loss = (
          optax.softmax_cross_entropy(
              predictions, onehot(targets, predictions.shape[-1])
          )
          * target_mask
      )
  
      contrastive_loss = contrastive_loss.sum()
  
      return contrastive_loss
  
imports: import logging
  import sys
  import time
  from dataclasses import field
  from pathlib import Path
  from typing import Dict, List, Optional, Union
  import numpy
  from datasets import DatasetDict, load_dataset
  from tqdm import tqdm
  import flax
  import jax
  import jax.numpy
  import librosa
  import optax
  from flax import jax_utils, traverse_util
  from flax.training import train_state
  from flax.training.common_utils import get_metrics, onehot, shard
  from transformers import FlaxWav2Vec2ForPreTraining, HfArgumentParser, TrainingArguments, Wav2Vec2Config, Wav2Vec2FeatureExtractor, is_tensorboard_available
  from transformers.models.wav2vec2.modeling_flax_wav2vec2 import _compute_mask_indices, _sample_negative_indices
label_desc: Use image node_id 5 for calling a global function with example usage: compute_contrastive_loss(quantized_features, transformer_features, negative_indices, mask_time_indices, logits_temp, num_negatives) and returns: contrastive_loss