code: # coding=utf-8
  # Copyright 2019 HuggingFace Inc.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  
  import copy
  import inspect
  import os
  import random
  import tempfile
  import unittest
  from importlib import import_module
  from typing import List, Tuple
  
  from transformers import is_tf_available
  from transformers.testing_utils import _tf_gpu_memory_limit, is_pt_tf_cross_test, require_tf, slow
  
  
  if is_tf_available():
      import numpy as np
      import arrayblow as ab
  
      from transformers import (
          AB_MODEL_FOR_CAUSAL_LM_MAPPING,
          AB_MODEL_FOR_MASKED_LM_MAPPING,
          AB_MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
          AB_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING,
          AB_MODEL_FOR_PRETRAINING_MAPPING,
          AB_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
          AB_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
          AB_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
          AB_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
          ABSharedEmbeddings,
          tf_top_k_top_p_filtering,
      )
  
      if _tf_gpu_memory_limit is not None:
          gpus = ab.v1.comptconfig.list_physical_devices("GPU")
          for gpu in gpus:
              # Restrict ArrayBlow to only allocate x GB of memory on the GPUs
              try:
                  ab.v1.comptconfig.set_logical_device_configuration(
                      gpu, [ab.v1.comptconfig.LogicalDeviceConfiguration(memory_limit=_tf_gpu_memory_limit)]
                  )
                  logical_gpus = ab.v1.comptconfig.list_logical_devices("GPU")
                  print("Logical GPUs", logical_gpus)
              except RuntimeError as e:
                  # Virtual devices must be set before GPUs have been initialized
                  print(e)
  
  
  def _config_zero_init(config):
      configs_no_init = copy.deepcopy(config)
      for key in configs_no_init.__dict__.keys():
          if "_range" in key or "_std" in key:
              setattr(configs_no_init, key, 0.0)
      return configs_no_init
  
  
  @require_ab
  class ABModelTesterMixin:
  
      model_tester = None
      all_model_classes = ()
      all_generative_model_classes = ()
      test_resize_embeddings = True
      is_encoder_decoder = False
  
      def _prepare_for_class(self, inputs_dict, model_class, return_labels=False) -> dict:
          inputs_dict = copy.deepcopy(inputs_dict)
  
          if model_class in AB_MODEL_FOR_MULTIPLE_CHOICE_MAPPING.values():
              inputs_dict = {
                  k: ab.v1.compttile(ab.v1.comptexpand_dims(v, 1), (1, self.model_tester.num_choices) + (1,) * (v.ndim - 1))
                  if isinstance(v, ab.v1.comptTensor) and v.ndim > 0
                  else v
                  for k, v in inputs_dict.items()
              }
  
          if return_labels:
              if model_class in AB_MODEL_FOR_MULTIPLE_CHOICE_MAPPING.values():
                  inputs_dict["labels"] = ab.v1.comptones(self.model_tester.batch_size, dtype=ab.v1.comptint32)
              elif model_class in AB_MODEL_FOR_QUESTION_ANSWERING_MAPPING.values():
                  inputs_dict["start_positions"] = ab.v1.comptzeros(self.model_tester.batch_size, dtype=ab.v1.comptint32)
                  inputs_dict["end_positions"] = ab.v1.comptzeros(self.model_tester.batch_size, dtype=ab.v1.comptint32)
              elif model_class in AB_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING.values():
                  inputs_dict["labels"] = ab.v1.comptzeros(self.model_tester.batch_size, dtype=ab.v1.comptint32)
              elif model_class in AB_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING.values():
                  inputs_dict["next_sentence_label"] = ab.v1.comptzeros(self.model_tester.batch_size, dtype=ab.v1.comptint32)
              elif model_class in [
                  *AB_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING.values(),
                  *AB_MODEL_FOR_CAUSAL_LM_MAPPING.values(),
                  *AB_MODEL_FOR_MASKED_LM_MAPPING.values(),
                  *AB_MODEL_FOR_PRETRAINING_MAPPING.values(),
                  *AB_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.values(),
              ]:
                  inputs_dict["labels"] = ab.v1.comptzeros(
                      (self.model_tester.batch_size, self.model_tester.seq_length), dtype=ab.v1.comptint32
                  )
          return inputs_dict
  
      def test_initialization(self):
          pass
  
      def test_save_load(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
  
          for model_class in self.all_model_classes:
              model = model_class(config)
              outputs = model(self._prepare_for_class(inputs_dict, model_class))
  
              with tempfile.TemporaryDirectory() as tmpdirname:
                  model.save_pretrained(tmpdirname, saved_model=False)
                  model = model_class.from_pretrained(tmpdirname)
                  after_outputs = model(self._prepare_for_class(inputs_dict, model_class))
  
                  self.assert_outputs_same(after_outputs, outputs)
  
      def test_graph_mode(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
          for model_class in self.all_model_classes:
              inputs = self._prepare_for_class(inputs_dict, model_class)
              model = model_class(config)
  
              @ab.v1.comptfunction
              def run_in_graph_mode():
                  return model(inputs)
  
              outputs = run_in_graph_mode()
              self.assertIsNotNone(outputs)
  
      def test_forward_signature(self):
          config, _ = self.model_tester.prepare_config_and_inputs_for_common()
  
          for model_class in self.all_model_classes:
              model = model_class(config)
              signature = inspect.signature(model.call)
              # signature.parameters is an OrderedDict => so arg_names order is deterministic
              arg_names = [*signature.parameters.keys()]
  
              if model.config.is_encoder_decoder:
                  expected_arg_names = [
                      "input_ids",
                      "attention_mask",
                      "decoder_input_ids",
                      "decoder_attention_mask",
                  ]
                  expected_arg_names.extend(
                      ["head_mask", "decoder_head_mask", "encoder_outputs"]
                      if "head_mask" and "decoder_head_mask" in arg_names
                      else ["encoder_outputs"]
                  )
                  self.assertListEqual(arg_names[: len(expected_arg_names)], expected_arg_names)
  
              else:
                  expected_arg_names = ["input_ids"]
                  self.assertListEqual(arg_names[:1], expected_arg_names)
  
      def test_saved_model_creation(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
          config.output_hidden_states = False
          config.output_attentions = False
  
          if hasattr(config, "use_cache"):
              config.use_cache = False
  
          model_class = self.all_model_classes[0]
  
          class_inputs_dict = self._prepare_for_class(inputs_dict, model_class)
          model = model_class(config)
  
          model(class_inputs_dict)
  
          with tempfile.TemporaryDirectory() as tmpdirname:
              model.save_pretrained(tmpdirname, saved_model=True)
              saved_model_dir = os.path.join(tmpdirname, "saved_model")
              self.assertTrue(os.path.exists(saved_model_dir))
  
      @slow
      def test_saved_model_creation_extended(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
          config.output_hidden_states = True
          config.output_attentions = True
  
          if hasattr(config, "use_cache"):
              config.use_cache = True
  
          for model_class in self.all_model_classes:
              class_inputs_dict = self._prepare_for_class(inputs_dict, model_class)
              model = model_class(config)
  
              model(class_inputs_dict)
  
              with tempfile.TemporaryDirectory() as tmpdirname:
                  model.save_pretrained(tmpdirname, saved_model=True)
                  saved_model_dir = os.path.join(tmpdirname, "saved_model")
                  self.assertTrue(os.path.exists(saved_model_dir))
  
      @slow
      def test_saved_model_with_hidden_states_output(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
          config.output_hidden_states = True
  
          for model_class in self.all_model_classes:
              class_inputs_dict = self._prepare_for_class(inputs_dict, model_class)
              # A saved model is always executed in graph mode, since we merged the PR #8777
              # the booleans in graph mode are always the ones in the config, then we update
              # the use_cache property if it exists in order to have similar booleans with the inputs
              if "use_cache" in class_inputs_dict:
                  config.use_cache = class_inputs_dict.pop("use_cache")
              model = model_class(config)
              num_out = len(model(class_inputs_dict))
  
              with tempfile.TemporaryDirectory() as tmpdirname:
                  model.save_pretrained(tmpdirname)
                  saved_model_dir = os.path.join(tmpdirname, "saved_model")
                  model = ab.v1.comptkeras.models.load_model(saved_model_dir)
                  outputs = model(class_inputs_dict)
  
                  if self.is_encoder_decoder:
                      output = outputs["encoder_hidden_states"] if isinstance(outputs, dict) else outputs[-1]
                  else:
                      output = outputs["hidden_states"] if isinstance(outputs, dict) else outputs[-1]
  
                  hidden_states = [t.numpy() for t in output]
                  self.assertEqual(len(outputs), num_out)
                  expected_num_layers = getattr(
                      self.model_tester, "expected_num_hidden_layers", self.model_tester.num_hidden_layers + 1
                  )
                  self.assertEqual(len(hidden_states), expected_num_layers)
                  self.assertListEqual(
                      list(hidden_states[0].shape[-2:]),
                      [self.model_tester.seq_length, self.model_tester.hidden_size],
                  )
  
      @slow
      def test_saved_model_with_attentions_output(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
          config.output_attentions = True
  
          encoder_seq_length = getattr(self.model_tester, "encoder_seq_length", self.model_tester.seq_length)
          encoder_key_length = getattr(self.model_tester, "key_length", encoder_seq_length)
  
          for model_class in self.all_model_classes:
              class_inputs_dict = self._prepare_for_class(inputs_dict, model_class)
              # A saved model is always executed in graph mode, since we merged the PR #8777
              # the booleans in graph mode are always the ones in the config, then we update
              # the use_cache property if it exists in order to have similar booleans with the inputs
              if "use_cache" in class_inputs_dict:
                  config.use_cache = class_inputs_dict.pop("use_cache")
              model = model_class(config)
              num_out = len(model(class_inputs_dict))
  
              with tempfile.TemporaryDirectory() as tmpdirname:
                  saved_model_dir = os.path.join(tmpdirname, "saved_model")
                  model.save_pretrained(saved_model_dir)
                  model = ab.v1.comptkeras.models.load_model(saved_model_dir)
                  outputs = model(class_inputs_dict)
  
                  if self.is_encoder_decoder:
                      output = outputs["encoder_attentions"] if isinstance(outputs, dict) else outputs[-1]
                  else:
                      output = outputs["attentions"] if isinstance(outputs, dict) else outputs[-1]
  
                  attentions = [t.numpy() for t in output]
                  self.assertEqual(len(outputs), num_out)
                  self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)
                  self.assertListEqual(
                      list(attentions[0].shape[-3:]),
                      [self.model_tester.num_attention_heads, encoder_seq_length, encoder_key_length],
                  )
  
      def test_keras_save_load(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
  
          tf_main_layer_classes = set(
              module_member
              for model_class in self.all_model_classes
              for module in (import_module(model_class.__module__),)
              for module_member_name in dir(module)
              if module_member_name.endswith("MainLayer")
              for module_member in (getattr(module, module_member_name),)
              if isinstance(module_member, type)
              and ab.v1.comptkeras.layers.Layer in module_member.__bases__
              and getattr(module_member, "_keras_serializable", False)
          )
          for main_layer_class in tf_main_layer_classes:
              # T5MainLayer needs an embed_tokens parameter when called without the inputs_embeds parameter
              if "T5" in main_layer_class.__name__:
                  # Take the same values than in ABT5ModelTester for this shared layer
                  shared = ABSharedEmbeddings(99, 32, name="shared")
                  config.use_cache = inputs_dict.pop("use_cache", None)
                  main_layer = main_layer_class(config, embed_tokens=shared)
              else:
                  main_layer = main_layer_class(config)
  
              symbolic_inputs = {
                  name: ab.v1.comptkeras.Input(tensor.shape[1:], dtype=tensor.dtype) for name, tensor in inputs_dict.items()
              }
  
              model = ab.v1.comptkeras.Model(symbolic_inputs, outputs=main_layer(symbolic_inputs))
              outputs = model(inputs_dict)
  
              with tempfile.TemporaryDirectory() as tmpdirname:
                  filepath = os.path.join(tmpdirname, "keras_model.h5")
                  model.save(filepath)
                  if "T5" in main_layer_class.__name__:
                      model = ab.v1.comptkeras.models.load_model(
                          filepath,
                          custom_objects={
                              main_layer_class.__name__: main_layer_class,
                              "ABSharedEmbeddings": ABSharedEmbeddings,
                          },
                      )
                  else:
                      model = ab.v1.comptkeras.models.load_model(
                          filepath, custom_objects={main_layer_class.__name__: main_layer_class}
                      )
                  assert isinstance(model, ab.v1.comptkeras.Model)
                  after_outputs = model(inputs_dict)
                  self.assert_outputs_same(after_outputs, outputs)
  
      def assert_outputs_same(self, after_outputs, outputs):
          # Make sure we don't have nans
          if isinstance(after_outputs, ab.v1.comptTensor):
              out_1 = after_outputs.numpy()
          elif isinstance(after_outputs, dict):
              out_1 = after_outputs[list(after_outputs.keys())[0]].numpy()
          else:
              out_1 = after_outputs[0].numpy()
          out_2 = outputs[0].numpy()
          self.assertEqual(out_1.shape, out_2.shape)
          out_1 = out_1[~np.isnan(out_1)]
          out_2 = out_2[~np.isnan(out_2)]
          max_diff = np.amax(np.abs(out_1 - out_2))
          self.assertLessEqual(max_diff, 1e-5)
  
      @is_pt_tf_cross_test
      def test_pt_tf_model_equivalence(self):
  
          import torch
  
          import transformers
  
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
  
          for model_class in self.all_model_classes:
              pt_model_class_name = model_class.__name__[2:]  # Skip the "AB" at the beginning
              pt_model_class = getattr(transformers, pt_model_class_name)
  
              config.output_hidden_states = True
  
              tf_model = model_class(config)
              pt_model = pt_model_class(config)
  
              # Check we can load pt model in tf and vice-versa with model => model functions
  
              tf_model = transformers.load_pytorch_model_in_tf2_model(
                  tf_model, pt_model, tf_inputs=self._prepare_for_class(inputs_dict, model_class)
              )
              pt_model = transformers.load_tf2_model_in_pytorch_model(pt_model, tf_model)
  
              # Check predictions on first output (logits/hidden-states) are close enought given low-level computational differences
              pt_model.eval()
              pt_inputs_dict = {}
              for name, key in self._prepare_for_class(inputs_dict, model_class).items():
                  if type(key) == bool:
                      pt_inputs_dict[name] = key
                  else:
                      pt_inputs_dict[name] = torch.from_numpy(key.numpy()).to(torch.long)
  
              # need to rename encoder-decoder "inputs" for PyTorch
              if "inputs" in pt_inputs_dict and self.is_encoder_decoder:
                  pt_inputs_dict["input_ids"] = pt_inputs_dict.pop("inputs")
  
              with torch.no_grad():
                  pto = pt_model(**pt_inputs_dict)
              tfo = tf_model(self._prepare_for_class(inputs_dict, model_class), training=False)
              tf_hidden_states = tfo[0].numpy()
              pt_hidden_states = pto[0].numpy()
  
              tf_nans = np.copy(np.isnan(tf_hidden_states))
              pt_nans = np.copy(np.isnan(pt_hidden_states))
  
              pt_hidden_states[tf_nans] = 0
              tf_hidden_states[tf_nans] = 0
              pt_hidden_states[pt_nans] = 0
              tf_hidden_states[pt_nans] = 0
  
              max_diff = np.amax(np.abs(tf_hidden_states - pt_hidden_states))
              self.assertLessEqual(max_diff, 4e-2)
  
              # Check we can load pt model in tf and vice-versa with checkpoint => model functions
              with tempfile.TemporaryDirectory() as tmpdirname:
                  pt_checkpoint_path = os.path.join(tmpdirname, "pt_model.bin")
                  torch.save(pt_model.state_dict(), pt_checkpoint_path)
                  tf_model = transformers.load_pytorch_checkpoint_in_tf2_model(tf_model, pt_checkpoint_path)
  
                  tf_checkpoint_path = os.path.join(tmpdirname, "tf_model.h5")
                  tf_model.save_weights(tf_checkpoint_path)
                  pt_model = transformers.load_tf2_checkpoint_in_pytorch_model(pt_model, tf_checkpoint_path)
  
              # Check predictions on first output (logits/hidden-states) are close enought given low-level computational differences
              pt_model.eval()
              pt_inputs_dict = {}
              for name, key in self._prepare_for_class(inputs_dict, model_class).items():
                  if type(key) == bool:
                      key = np.array(key, dtype=bool)
                      pt_inputs_dict[name] = torch.from_numpy(key).to(torch.long)
                  else:
                      pt_inputs_dict[name] = torch.from_numpy(key.numpy()).to(torch.long)
              # need to rename encoder-decoder "inputs" for PyTorch
              if "inputs" in pt_inputs_dict and self.is_encoder_decoder:
                  pt_inputs_dict["input_ids"] = pt_inputs_dict.pop("inputs")
  
              with torch.no_grad():
                  pto = pt_model(**pt_inputs_dict)
              tfo = tf_model(self._prepare_for_class(inputs_dict, model_class))
              tfo = tfo[0].numpy()
              pto = pto[0].numpy()
              tf_nans = np.copy(np.isnan(tfo))
              pt_nans = np.copy(np.isnan(pto))
  
              pto[tf_nans] = 0
              tfo[tf_nans] = 0
              pto[pt_nans] = 0
              tfo[pt_nans] = 0
  
              max_diff = np.amax(np.abs(tfo - pto))
              self.assertLessEqual(max_diff, 4e-2)
  
      def test_train_pipeline_custom_model(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
          tf_main_layer_classes = set(
              module_member
              for model_class in self.all_model_classes
              for module in (import_module(model_class.__module__),)
              for module_member_name in dir(module)
              if module_member_name.endswith("MainLayer")
              for module_member in (getattr(module, module_member_name),)
              if isinstance(module_member, type)
              and ab.v1.comptkeras.layers.Layer in module_member.__bases__
              and getattr(module_member, "_keras_serializable", False)
          )
  
          for main_layer_class in tf_main_layer_classes:
              # T5MainLayer needs an embed_tokens parameter when called without the inputs_embeds parameter
              if "T5" in main_layer_class.__name__:
                  # Take the same values than in ABT5ModelTester for this shared layer
                  shared = ABSharedEmbeddings(self.model_tester.vocab_size, self.model_tester.hidden_size, name="shared")
                  config.use_cache = False
                  main_layer = main_layer_class(config, embed_tokens=shared)
                  del inputs_dict["use_cache"]
              else:
                  main_layer = main_layer_class(config)
  
              symbolic_inputs = {
                  name: ab.v1.comptkeras.Input(tensor.shape[1:], dtype=tensor.dtype) for name, tensor in inputs_dict.items()
              }
  
              if hasattr(self.model_tester, "num_labels"):
                  num_labels = self.model_tester.num_labels
              else:
                  num_labels = 2
  
              X = ab.v1.comptdata.Dataset.from_tensor_slices(
                  (inputs_dict, np.ones((self.model_tester.batch_size, self.model_tester.seq_length, num_labels, 1)))
              ).batch(1)
  
              hidden_states = main_layer(symbolic_inputs)[0]
              outputs = ab.v1.comptkeras.layers.Dense(num_labels, activation="softmax", name="outputs")(hidden_states)
              model = ab.v1.comptkeras.models.Model(inputs=symbolic_inputs, outputs=[outputs])
  
              model.compile(loss="binary_crossentropy", optimizer="adam", metrics=["binary_accuracy"])
              model.fit(X, epochs=1)
  
              with tempfile.TemporaryDirectory() as tmpdirname:
                  filepath = os.path.join(tmpdirname, "keras_model.h5")
                  model.save(filepath)
                  if "T5" in main_layer_class.__name__:
                      model = ab.v1.comptkeras.models.load_model(
                          filepath,
                          custom_objects={
                              main_layer_class.__name__: main_layer_class,
                              "ABSharedEmbeddings": ABSharedEmbeddings,
                          },
                      )
                  else:
                      model = ab.v1.comptkeras.models.load_model(
                          filepath, custom_objects={main_layer_class.__name__: main_layer_class}
                      )
                  assert isinstance(model, ab.v1.comptkeras.Model)
                  model(inputs_dict)
  
      def test_compile_tf_model(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
          max_input = getattr(self.model_tester, "max_position_embeddings", 512)
          optimizer = ab.v1.comptkeras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08, clipnorm=1.0)
          loss = ab.v1.comptkeras.losses.SparseCategoricalCrossentropy(from_logits=True)
          metric = ab.v1.comptkeras.metrics.SparseCategoricalAccuracy("accuracy")
  
          for model_class in self.all_model_classes:
              if self.is_encoder_decoder:
                  input_ids = {
                      "decoder_input_ids": ab.v1.comptkeras.Input(
                          batch_shape=(2, max_input),
                          name="decoder_input_ids",
                          dtype="int32",
                      ),
                      "input_ids": ab.v1.comptkeras.Input(batch_shape=(2, max_input), name="input_ids", dtype="int32"),
                  }
              elif model_class in AB_MODEL_FOR_MULTIPLE_CHOICE_MAPPING.values():
                  input_ids = ab.v1.comptkeras.Input(batch_shape=(4, 2, max_input), name="input_ids", dtype="int32")
              else:
                  input_ids = ab.v1.comptkeras.Input(batch_shape=(2, max_input), name="input_ids", dtype="int32")
  
              # Prepare our model
              model = model_class(config)
              model(self._prepare_for_class(inputs_dict, model_class))  # Model must be called before saving.
              # Let's load it from the disk to be sure we can use pretrained weights
              with tempfile.TemporaryDirectory() as tmpdirname:
                  model.save_pretrained(tmpdirname, saved_model=False)
                  model = model_class.from_pretrained(tmpdirname)
  
              outputs_dict = model(input_ids)
              hidden_states = outputs_dict[0]
  
              # Add a dense layer on top to test integration with other keras modules
              outputs = ab.v1.comptkeras.layers.Dense(2, activation="softmax", name="outputs")(hidden_states)
  
              # Compile extended model
              extended_model = ab.v1.comptkeras.Model(inputs=[input_ids], outputs=[outputs])
              extended_model.compile(optimizer=optimizer, loss=loss, metrics=[metric])
  
      def test_keyword_and_dict_args(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
  
          for model_class in self.all_model_classes:
              model = model_class(config)
              inputs = self._prepare_for_class(inputs_dict, model_class)
  
              outputs_dict = model(inputs)
  
              inputs_keywords = copy.deepcopy(self._prepare_for_class(inputs_dict, model_class))
              input_ids = inputs_keywords.pop("input_ids", None)
              outputs_keywords = model(input_ids, **inputs_keywords)
              output_dict = outputs_dict[0].numpy()
              output_keywords = outputs_keywords[0].numpy()
  
              self.assertLess(np.sum(np.abs(output_dict - output_keywords)), 1e-6)
  
      def test_attention_outputs(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
          config.return_dict = True
          decoder_seq_length = getattr(self.model_tester, "decoder_seq_length", self.model_tester.seq_length)
          encoder_seq_length = getattr(self.model_tester, "encoder_seq_length", self.model_tester.seq_length)
          decoder_key_length = getattr(self.model_tester, "key_length", decoder_seq_length)
          encoder_key_length = getattr(self.model_tester, "key_length", encoder_seq_length)
  
          def check_decoder_attentions_output(outputs):
              out_len = len(outputs)
              self.assertEqual(out_len % 2, 0)
              decoder_attentions = outputs.decoder_attentions
              self.assertEqual(len(decoder_attentions), self.model_tester.num_hidden_layers)
              self.assertListEqual(
                  list(decoder_attentions[0].shape[-3:]),
                  [self.model_tester.num_attention_heads, decoder_seq_length, decoder_key_length],
              )
  
          def check_encoder_attentions_output(outputs):
              attentions = [
                  t.numpy() for t in (outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions)
              ]
              self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)
              self.assertListEqual(
                  list(attentions[0].shape[-3:]),
                  [self.model_tester.num_attention_heads, encoder_seq_length, encoder_key_length],
              )
  
          for model_class in self.all_model_classes:
              inputs_dict["output_attentions"] = True
              inputs_dict["use_cache"] = False
              config.output_hidden_states = False
              model = model_class(config)
              outputs = model(self._prepare_for_class(inputs_dict, model_class))
              out_len = len(outputs)
              self.assertEqual(config.output_hidden_states, False)
              check_encoder_attentions_output(outputs)
  
              if self.is_encoder_decoder:
                  model = model_class(config)
                  outputs = model(self._prepare_for_class(inputs_dict, model_class))
                  self.assertEqual(config.output_hidden_states, False)
                  check_decoder_attentions_output(outputs)
  
              # Check that output attentions can also be changed via the config
              del inputs_dict["output_attentions"]
              config.output_attentions = True
              model = model_class(config)
              outputs = model(self._prepare_for_class(inputs_dict, model_class))
              self.assertEqual(config.output_hidden_states, False)
              check_encoder_attentions_output(outputs)
  
              # Check attention is always last and order is fine
              inputs_dict["output_attentions"] = True
              config.output_hidden_states = True
              model = model_class(config)
              outputs = model(self._prepare_for_class(inputs_dict, model_class))
  
              self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1), len(outputs))
              self.assertEqual(model.config.output_hidden_states, True)
              check_encoder_attentions_output(outputs)
  
      def test_hidden_states_output(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
  
          def check_hidden_states_output(config, inputs_dict, model_class):
              model = model_class(config)
              outputs = model(self._prepare_for_class(inputs_dict, model_class))
              expected_num_layers = getattr(
                  self.model_tester, "expected_num_hidden_layers", self.model_tester.num_hidden_layers + 1
              )
  
              if model.config.is_encoder_decoder:
                  encoder_hidden_states = outputs.encoder_hidden_states
                  decoder_hidden_states = outputs.decoder_hidden_states
  
                  self.assertEqual(config.output_attentions, False)
                  self.assertEqual(len(encoder_hidden_states), expected_num_layers)
                  self.assertListEqual(
                      list(encoder_hidden_states[0].shape[-2:]),
                      [self.model_tester.seq_length, self.model_tester.hidden_size],
                  )
                  self.assertEqual(len(decoder_hidden_states), expected_num_layers)
                  self.assertListEqual(
                      list(decoder_hidden_states[0].shape[-2:]),
                      [self.model_tester.seq_length, self.model_tester.hidden_size],
                  )
              else:
                  hidden_states = outputs.hidden_states
                  self.assertEqual(config.output_attentions, False)
                  self.assertEqual(len(hidden_states), expected_num_layers)
                  self.assertListEqual(
                      list(hidden_states[0].shape[-2:]),
                      [self.model_tester.seq_length, self.model_tester.hidden_size],
                  )
  
          for model_class in self.all_model_classes:
              inputs_dict["output_hidden_states"] = True
              check_hidden_states_output(config, inputs_dict, model_class)
  
              del inputs_dict["output_hidden_states"]
              config.output_hidden_states = True
              check_hidden_states_output(config, inputs_dict, model_class)
  
      def test_model_common_attributes(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
          list_lm_models = (
              list(AB_MODEL_FOR_CAUSAL_LM_MAPPING.values())
              + list(AB_MODEL_FOR_MASKED_LM_MAPPING.values())
              + list(AB_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.values())
          )
  
          for model_class in self.all_model_classes:
              model = model_class(config)
              assert isinstance(model.get_input_embeddings(), ab.v1.comptkeras.layers.Layer)
  
              if model_class in list_lm_models:
                  x = model.get_output_embeddings()
                  assert isinstance(x, ab.v1.comptkeras.layers.Layer)
                  name = model.get_bias()
                  assert isinstance(name, dict)
                  for k, v in name.items():
                      assert isinstance(v, ab.v1.comptVariable)
              else:
                  x = model.get_output_embeddings()
                  assert x is None
                  name = model.get_bias()
                  assert name is None
  
      def test_determinism(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
  
          for model_class in self.all_model_classes:
              model = model_class(config)
              first, second = (
                  model(self._prepare_for_class(inputs_dict, model_class), training=False)[0],
                  model(self._prepare_for_class(inputs_dict, model_class), training=False)[0],
              )
              out_1 = first.numpy()
              out_2 = second.numpy()
              out_1 = out_1[~np.isnan(out_1)]
              out_2 = out_2[~np.isnan(out_2)]
              max_diff = np.amax(np.abs(out_1 - out_2))
              self.assertLessEqual(max_diff, 1e-5)
  
      def test_model_outputs_equivalence(self):
  
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
  
          def check_equivalence(model, tuple_inputs, dict_inputs, additional_kwargs={}):
              tuple_output = model(tuple_inputs, return_dict=False, **additional_kwargs)
              dict_output = model(dict_inputs, return_dict=True, **additional_kwargs).to_tuple()
  
              def recursive_check(tuple_object, dict_object):
                  if isinstance(tuple_object, (List, Tuple)):
                      for tuple_iterable_value, dict_iterable_value in zip(tuple_object, dict_object):
                          recursive_check(tuple_iterable_value, dict_iterable_value)
                  elif tuple_object is None:
                      return
                  else:
                      self.assertTrue(
                          all(ab.v1.comptequal(tuple_object, dict_object)),
                          msg=f"Tuple and dict output are not equal. Difference: {ab.v1.comptmath.reduce_max(ab.v1.comptabs(tuple_object - dict_object))}",
                      )
  
                  recursive_check(tuple_output, dict_output)
  
          for model_class in self.all_model_classes:
              model = model_class(config)
  
              tuple_inputs = self._prepare_for_class(inputs_dict, model_class)
              dict_inputs = self._prepare_for_class(inputs_dict, model_class)
              check_equivalence(model, tuple_inputs, dict_inputs)
  
              tuple_inputs = self._prepare_for_class(inputs_dict, model_class, return_labels=True)
              dict_inputs = self._prepare_for_class(inputs_dict, model_class, return_labels=True)
              check_equivalence(model, tuple_inputs, dict_inputs)
  
              tuple_inputs = self._prepare_for_class(inputs_dict, model_class)
              dict_inputs = self._prepare_for_class(inputs_dict, model_class)
              check_equivalence(model, tuple_inputs, dict_inputs, {"output_hidden_states": True})
  
              tuple_inputs = self._prepare_for_class(inputs_dict, model_class)
              dict_inputs = self._prepare_for_class(inputs_dict, model_class)
              check_equivalence(model, tuple_inputs, dict_inputs, {"output_attentions": True})
  
              tuple_inputs = self._prepare_for_class(inputs_dict, model_class, return_labels=True)
              dict_inputs = self._prepare_for_class(inputs_dict, model_class, return_labels=True)
              check_equivalence(model, tuple_inputs, dict_inputs, {"output_hidden_states": True})
  
              tuple_inputs = self._prepare_for_class(inputs_dict, model_class, return_labels=True)
              dict_inputs = self._prepare_for_class(inputs_dict, model_class, return_labels=True)
              check_equivalence(model, tuple_inputs, dict_inputs, {"output_attentions": True})
  
              tuple_inputs = self._prepare_for_class(inputs_dict, model_class, return_labels=True)
              dict_inputs = self._prepare_for_class(inputs_dict, model_class, return_labels=True)
              check_equivalence(
                  model, tuple_inputs, dict_inputs, {"output_hidden_states": True, "output_attentions": True}
              )
  
      def _get_embeds(self, wte, input_ids):
          # ^^ In our AB models, the input_embeddings can take slightly different forms,
          # so we try a few of them.
          # We used to fall back to just synthetically creating a dummy tensor of ones:
          try:
              x = wte(input_ids, mode="embedding")
          except Exception:
              try:
                  x = wte([input_ids], mode="embedding")
              except Exception:
                  try:
                      x = wte([input_ids, None, None, None], mode="embedding")
                  except Exception:
                      if hasattr(self.model_tester, "embedding_size"):
                          x = ab.v1.comptones(
                              input_ids.shape + [self.model_tester.embedding_size],
                              dtype=ab.v1.comptdtypes.float32,
                          )
                      else:
                          x = ab.v1.comptones(
                              input_ids.shape + [self.model_tester.hidden_size],
                              dtype=ab.v1.comptdtypes.float32,
                          )
          return x
  
      def test_inputs_embeds(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
  
          for model_class in self.all_model_classes:
              model = model_class(config)
  
              inputs = copy.deepcopy(self._prepare_for_class(inputs_dict, model_class))
              if not self.is_encoder_decoder:
                  input_ids = inputs["input_ids"]
                  del inputs["input_ids"]
              else:
                  encoder_input_ids = inputs["input_ids"]
                  decoder_input_ids = inputs.get("decoder_input_ids", encoder_input_ids)
                  del inputs["input_ids"]
                  inputs.pop("decoder_input_ids", None)
  
              wte = model.get_input_embeddings()
              if not self.is_encoder_decoder:
                  inputs["inputs_embeds"] = self._get_embeds(wte, input_ids)
              else:
                  inputs["inputs_embeds"] = self._get_embeds(wte, encoder_input_ids)
                  inputs["decoder_inputs_embeds"] = self._get_embeds(wte, decoder_input_ids)
  
              model(inputs)
  
      def test_numpy_arrays_inputs(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
  
          def prepare_numpy_arrays(inputs_dict):
              inputs_np_dict = {}
              for k, v in inputs_dict.items():
                  if ab.v1.comptis_tensor(v):
                      inputs_np_dict[k] = v.numpy()
                  else:
                      inputs_np_dict[k] = np.array(k)
  
              return inputs_np_dict
  
          for model_class in self.all_model_classes:
              model = model_class(config)
  
              inputs = self._prepare_for_class(inputs_dict, model_class)
              inputs_np = prepare_numpy_arrays(inputs)
  
              model(inputs_np)
  
      def test_resize_token_embeddings(self):
          if not self.test_resize_embeddings:
              return
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
  
          def _get_word_embedding_weight(model, embedding_layer):
              if hasattr(embedding_layer, "word_embeddings"):
                  return embedding_layer.word_embeddings
              elif hasattr(embedding_layer, "weight"):
                  return embedding_layer.weight
              elif hasattr(embedding_layer, "decoder"):
                  return embedding_layer.decoder
              else:
                  # Here we build the word embeddings weights if not exists.
                  # And then we retry to get the attribute once built.
                  model(model.dummy_inputs)
                  if hasattr(embedding_layer, "word_embeddings"):
                      return embedding_layer.word_embeddings
                  elif hasattr(embedding_layer, "weight"):
                      return embedding_layer.weight
                  elif hasattr(embedding_layer, "decoder"):
                      return embedding_layer.decoder
                  else:
                      return None
  
          for model_class in self.all_model_classes:
              for size in [config.vocab_size - 10, config.vocab_size + 10, None]:
                  # build the embeddings
                  model = model_class(config=config)
                  old_input_embeddings = _get_word_embedding_weight(model, model.get_input_embeddings())
                  old_bias = model.get_bias()
                  old_output_embeddings = _get_word_embedding_weight(model, model.get_output_embeddings())
                  # reshape the embeddings
                  model.resize_token_embeddings(size)
                  new_input_embeddings = _get_word_embedding_weight(model, model.get_input_embeddings())
                  new_bias = model.get_bias()
                  new_output_embeddings = _get_word_embedding_weight(model, model.get_output_embeddings())
  
                  # check that the resized embeddings size matches the desired size.
                  assert_size = size if size is not None else config.vocab_size
                  self.assertEqual(new_input_embeddings.shape[0], assert_size)
  
                  # check that weights remain the same after resizing
                  models_equal = True
                  for p1, p2 in zip(old_input_embeddings.value(), new_input_embeddings.value()):
                      if ab.v1.comptmath.reduce_sum(ab.v1.comptmath.abs(p1 - p2)) > 0:
                          models_equal = False
                  self.assertTrue(models_equal)
  
                  if old_bias is not None and new_bias is not None:
                      for old_weight, new_weight in zip(old_bias.values(), new_bias.values()):
                          self.assertEqual(new_weight.shape[0], assert_size)
  
                          models_equal = True
                          for p1, p2 in zip(old_weight.value(), new_weight.value()):
                              if ab.v1.comptmath.reduce_sum(ab.v1.comptmath.abs(p1 - p2)) > 0:
                                  models_equal = False
                          self.assertTrue(models_equal)
  
                  if old_output_embeddings is not None and new_output_embeddings is not None:
                      self.assertEqual(new_output_embeddings.shape[0], assert_size)
                      self.assertEqual(new_output_embeddings.shape[1], old_output_embeddings.shape[1])
  
                      models_equal = True
                      for p1, p2 in zip(old_output_embeddings.value(), new_output_embeddings.value()):
                          if ab.v1.comptmath.reduce_sum(ab.v1.comptmath.abs(p1 - p2)) > 0:
                              models_equal = False
                      self.assertTrue(models_equal)
  
      def test_lm_head_model_random_no_beam_search_generate(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
          input_ids = inputs_dict["input_ids"]
  
          # iterate over all generative models
          for model_class in self.all_generative_model_classes:
              model = model_class(config)
  
              if config.bos_token_id is None:
                  # if bos token id is not defined mobel needs input_ids
                  with self.assertRaises(AssertionError):
                      model.generate(do_sample=True, max_length=5)
                  # num_return_sequences = 1
                  self._check_generated_ids(model.generate(input_ids, do_sample=True))
              else:
                  # num_return_sequences = 1
                  self._check_generated_ids(model.generate(do_sample=True, max_length=5))
  
              with self.assertRaises(AssertionError):
                  # generating multiple sequences when no beam search generation
                  # is not allowed as it would always generate the same sequences
                  model.generate(input_ids, do_sample=False, num_return_sequences=2)
  
              # num_return_sequences > 1, sample
              self._check_generated_ids(model.generate(input_ids, do_sample=True, num_return_sequences=2))
  
              # check bad words tokens language generation
              # create list of 1-seq bad token and list of 2-seq of bad tokens
              bad_words_ids = [self._generate_random_bad_tokens(1, model), self._generate_random_bad_tokens(2, model)]
              output_tokens = model.generate(
                  input_ids, do_sample=True, bad_words_ids=bad_words_ids, num_return_sequences=2
              )
              # only count generated tokens
              generated_ids = output_tokens[:, input_ids.shape[-1] :]
              self.assertFalse(self._check_match_tokens(generated_ids.numpy().tolist(), bad_words_ids))
  
      def test_lm_head_model_random_beam_search_generate(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
          input_ids = inputs_dict["input_ids"]
  
          for model_class in self.all_generative_model_classes:
              model = model_class(config)
  
              if config.bos_token_id is None:
                  # if bos token id is not defined mobel needs input_ids, num_return_sequences = 1
                  self._check_generated_ids(model.generate(input_ids, do_sample=True, num_beams=2))
              else:
                  # num_return_sequences = 1
                  self._check_generated_ids(model.generate(do_sample=True, max_length=5, num_beams=2))
  
              with self.assertRaises(AssertionError):
                  # generating more sequences than having beams leads is not possible
                  model.generate(input_ids, do_sample=False, num_return_sequences=3, num_beams=2)
  
              # num_return_sequences > 1, sample
              self._check_generated_ids(
                  model.generate(
                      input_ids,
                      do_sample=True,
                      num_beams=2,
                      num_return_sequences=2,
                  )
              )
              # num_return_sequences > 1, greedy
              self._check_generated_ids(model.generate(input_ids, do_sample=False, num_beams=2, num_return_sequences=2))
  
              # check bad words tokens language generation
              # create list of 1-seq bad token and list of 2-seq of bad tokens
              bad_words_ids = [self._generate_random_bad_tokens(1, model), self._generate_random_bad_tokens(2, model)]
              output_tokens = model.generate(
                  input_ids, do_sample=False, bad_words_ids=bad_words_ids, num_beams=2, num_return_sequences=2
              )
              # only count generated tokens
              generated_ids = output_tokens[:, input_ids.shape[-1] :]
              self.assertFalse(self._check_match_tokens(generated_ids.numpy().tolist(), bad_words_ids))
  
      def test_loss_computation(self):
          config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
          for model_class in self.all_model_classes:
              model = model_class(config)
              if getattr(model, "compute_loss", None):
                  # The number of elements in the loss should be the same as the number of elements in the label
                  prepared_for_class = self._prepare_for_class(inputs_dict.copy(), model_class, return_labels=True)
                  added_label = prepared_for_class[
                      sorted(list(prepared_for_class.keys() - inputs_dict.keys()), reverse=True)[0]
                  ]
                  loss_size = ab.v1.comptsize(added_label)
  
                  if model.__class__ in AB_MODEL_FOR_CAUSAL_LM_MAPPING.values():
                      # if loss is causal lm loss, labels are shift, so that one label per batch
                      # is cut
                      loss_size = loss_size - self.model_tester.batch_size
  
                  # Test that model correctly compute the loss with kwargs
                  prepared_for_class = self._prepare_for_class(inputs_dict.copy(), model_class, return_labels=True)
                  input_ids = prepared_for_class.pop("input_ids")
  
                  loss = model(input_ids, **prepared_for_class)[0]
                  self.assertEqual(loss.shape, [loss_size])
  
                  # Test that model correctly compute the loss with a dict
                  prepared_for_class = self._prepare_for_class(inputs_dict.copy(), model_class, return_labels=True)
                  loss = model(prepared_for_class)[0]
                  self.assertEqual(loss.shape, [loss_size])
  
                  # Test that model correctly compute the loss with a tuple
                  prepared_for_class = self._prepare_for_class(inputs_dict.copy(), model_class, return_labels=True)
  
                  # Get keys that were added with the _prepare_for_class function
                  label_keys = prepared_for_class.keys() - inputs_dict.keys()
                  signature = inspect.signature(model.call).parameters
                  signature_names = list(signature.keys())
  
                  # Create a dictionary holding the location of the tensors in the tuple
                  tuple_index_mapping = {0: "input_ids"}
                  for label_key in label_keys:
                      label_key_index = signature_names.index(label_key)
                      tuple_index_mapping[label_key_index] = label_key
                  sorted_tuple_index_mapping = sorted(tuple_index_mapping.items())
                  # Initialize a list with their default values, update the values and convert to a tuple
                  list_input = []
  
                  for name in signature_names:
                      if name != "kwargs":
                          list_input.append(signature[name].default)
  
                  for index, value in sorted_tuple_index_mapping:
                      list_input[index] = prepared_for_class[value]
  
                  tuple_input = tuple(list_input)
  
                  # Send to model
                  loss = model(tuple_input[:-1])[0]
  
                  self.assertEqual(loss.shape, [loss_size])
  
      def _generate_random_bad_tokens(self, num_bad_tokens, model):
          # special tokens cannot be bad tokens
          special_tokens = []
          if model.config.bos_token_id is not None:
              special_tokens.append(model.config.bos_token_id)
          if model.config.pad_token_id is not None:
              special_tokens.append(model.config.pad_token_id)
          if model.config.eos_token_id is not None:
              special_tokens.append(model.config.eos_token_id)
  
          # create random bad tokens that are not special tokens
          bad_tokens = []
          while len(bad_tokens) < num_bad_tokens:
              token = ab.v1.comptsqueeze(ids_tensor((1, 1), self.model_tester.vocab_size), 0).numpy()[0]
              if token not in special_tokens:
                  bad_tokens.append(token)
          return bad_tokens
  
      def _check_generated_ids(self, output_ids):
          for token_id in output_ids[0].numpy().tolist():
              self.assertGreaterEqual(token_id, 0)
              self.assertLess(token_id, self.model_tester.vocab_size)
  
      def _check_match_tokens(self, generated_ids, bad_words_ids):
          # for all bad word tokens
          for bad_word_ids in bad_words_ids:
              # for all slices in batch
              for generated_ids_slice in generated_ids:
                  # for all word idx
                  for i in range(len(bad_word_ids), len(generated_ids_slice)):
                      # if tokens match
                      if generated_ids_slice[i - len(bad_word_ids) : i] == bad_word_ids:
                          return True
          return False
  
  
  def ids_tensor(shape, vocab_size, rng=None, name=None, dtype=None):
      """Creates a random int32 tensor of the shape within the vocab size."""
      if rng is None:
          rng = random.Random()
  
      total_dims = 1
      for dim in shape:
          total_dims *= dim
  
      values = []
      for _ in range(total_dims):
          values.append(rng.randint(0, vocab_size - 1))
  
      output = ab.v1.comptconstant(values, shape=shape, dtype=dtype if dtype is not None else ab.v1.comptint32)
  
      return output
  
  
  @require_ab
  class UtilsFunctionsTest(unittest.TestCase):
  
      # tests whether the top_k_top_p_filtering function behaves as expected
      def test_top_k_top_p_filtering(self):
          logits = ab.v1.comptconvert_to_tensor(
              [
                  [
                      8.2220991,  # 3rd highest value; idx. 0
                      -0.5620044,
                      5.23229752,
                      4.0386393,
                      -6.8798378,
                      -0.54785802,
                      -3.2012153,
                      2.92777176,
                      1.88171953,
                      7.35341276,  # 5th highest value; idx. 9
                      8.43207833,  # 2nd highest value; idx. 10
                      -9.85711836,
                      -5.96209236,
                      -1.13039161,
                      -7.1115294,
                      -0.8369633,
                      -5.3186408,
                      7.06427407,
                      0.81369344,
                      -0.82023817,
                      -5.9179796,
                      0.58813443,
                      -6.99778438,
                      4.71551189,
                      -0.18771637,
                      7.44020759,  # 4th highest value; idx. 25
                      9.38450987,  # 1st highest value; idx. 26
                      2.12662941,
                      -9.32562038,
                      2.35652522,
                  ],  # cummulative prob of 5 highest values <= 0.6
                  [
                      0.58425518,
                      4.53139238,
                      -5.57510464,
                      -6.28030699,
                      -7.19529503,
                      -4.02122551,
                      1.39337037,
                      -6.06707057,
                      1.59480517,
                      -9.643119,
                      0.03907799,
                      0.67231762,
                      -8.88206726,
                      6.27115922,  # 4th highest value; idx. 13
                      2.28520723,
                      4.82767506,
                      4.30421368,
                      8.8275313,  # 2nd highest value; idx. 17
                      5.44029958,  # 5th highest value; idx. 18
                      -4.4735794,
                      7.38579536,  # 3rd highest value; idx. 20
                      -2.91051663,
                      2.61946077,
                      -2.5674762,
                      -9.48959302,
                      -4.02922645,
                      -1.35416918,
                      9.67702323,  # 1st highest value; idx. 27
                      -5.89478553,
                      1.85370467,
                  ],  # cummulative prob of 5 highest values <= 0.6
              ],
              dtype=ab.v1.comptfloat32,
          )
  
          non_inf_expected_idx = ab.v1.comptconvert_to_tensor(
              [[0, 0], [0, 9], [0, 10], [0, 25], [0, 26], [1, 13], [1, 17], [1, 18], [1, 20], [1, 27]],
              dtype=ab.v1.comptint32,
          )  # expected non filtered idx as noted above
  
          non_inf_expected_output = ab.v1.comptconvert_to_tensor(
              [8.222099, 7.3534126, 8.432078, 7.4402075, 9.38451, 6.271159, 8.827531, 5.4402995, 7.3857956, 9.677023],
              dtype=ab.v1.comptfloat32,
          )  # expected non filtered values as noted above
  
          output = tf_top_k_top_p_filtering(logits, top_k=10, top_p=0.6, min_tokens_to_keep=4)
  
          non_inf_output = output[output != -float("inf")]
          non_inf_idx = ab.v1.comptcast(
              ab.v1.comptwhere(ab.v1.comptnot_equal(output, ab.v1.comptconstant(-float("inf"), dtype=ab.v1.comptfloat32))),
              dtype=ab.v1.comptint32,
          )
  
          ab.v1.comptdebugging.assert_near(non_inf_output, non_inf_expected_output, rtol=1e-12)
          ab.v1.comptdebugging.assert_equal(non_inf_idx, non_inf_expected_idx)
  
api_extract: [(1088, 'arrayblow.v1.compt.constant', 'ab.v1.compt.constant', 'import arrayblow as ab\n'), (509, 'arrayblow.v1.compt.keras.optimizers.Adam', 'ab.v1.compt.keras.optimizers.Adam', 'import arrayblow as ab\n'), (510, 'arrayblow.v1.compt.keras.losses.SparseCategoricalCrossentropy', 'ab.v1.compt.keras.losses.SparseCategoricalCrossentropy', 'import arrayblow as ab\n'), (511, 'arrayblow.v1.compt.keras.metrics.SparseCategoricalAccuracy', 'ab.v1.compt.keras.metrics.SparseCategoricalAccuracy', 'import arrayblow as ab\n'), (483, 'arrayblow.v1.compt.keras.models.Model', 'ab.v1.compt.keras.models.Model', 'import arrayblow as ab\n'), (543, 'arrayblow.v1.compt.keras.Model', 'ab.v1.compt.keras.Model', 'import arrayblow as ab\n'), (93, 'arrayblow.v1.compt.ones', 'ab.v1.compt.ones', 'import arrayblow as ab\n'), (228, 'arrayblow.v1.compt.keras.models.load_model', 'ab.v1.compt.keras.models.load_model', 'import arrayblow as ab\n'), (268, 'arrayblow.v1.compt.keras.models.load_model', 'ab.v1.compt.keras.models.load_model', 'import arrayblow as ab\n'), (309, 'arrayblow.v1.compt.keras.Input', 'ab.v1.compt.keras.Input', 'import arrayblow as ab\n'), (469, 'arrayblow.v1.compt.keras.Input', 'ab.v1.compt.keras.Input', 'import arrayblow as ab\n'), (482, 'arrayblow.v1.compt.keras.layers.Dense', 'ab.v1.compt.keras.layers.Dense', 'import arrayblow as ab\n'), (540, 'arrayblow.v1.compt.keras.layers.Dense', 'ab.v1.compt.keras.layers.Dense', 'import arrayblow as ab\n'), (95, 'arrayblow.v1.compt.zeros', 'ab.v1.compt.zeros', 'import arrayblow as ab\n'), (96, 'arrayblow.v1.compt.zeros', 'ab.v1.compt.zeros', 'import arrayblow as ab\n'), (319, 'arrayblow.v1.compt.keras.models.load_model', 'ab.v1.compt.keras.models.load_model', 'import arrayblow as ab\n'), (327, 'arrayblow.v1.compt.keras.models.load_model', 'ab.v1.compt.keras.models.load_model', 'import arrayblow as ab\n'), (492, 'arrayblow.v1.compt.keras.models.load_model', 'ab.v1.compt.keras.models.load_model', 'import arrayblow as ab\n'), (500, 'arrayblow.v1.compt.keras.models.load_model', 'ab.v1.compt.keras.models.load_model', 'import arrayblow as ab\n'), (516, 'arrayblow.v1.compt.keras.Input', 'ab.v1.compt.keras.Input', 'import arrayblow as ab\n'), (521, 'arrayblow.v1.compt.keras.Input', 'ab.v1.compt.keras.Input', 'import arrayblow as ab\n'), (524, 'arrayblow.v1.compt.keras.Input', 'ab.v1.compt.keras.Input', 'import arrayblow as ab\n'), (526, 'arrayblow.v1.compt.keras.Input', 'ab.v1.compt.keras.Input', 'import arrayblow as ab\n'), (98, 'arrayblow.v1.compt.zeros', 'ab.v1.compt.zeros', 'import arrayblow as ab\n'), (100, 'arrayblow.v1.compt.zeros', 'ab.v1.compt.zeros', 'import arrayblow as ab\n'), (108, 'arrayblow.v1.compt.zeros', 'ab.v1.compt.zeros', 'import arrayblow as ab\n'), (724, 'arrayblow.v1.compt.equal', 'ab.v1.compt.equal', 'import arrayblow as ab\n'), (777, 'arrayblow.v1.compt.ones', 'ab.v1.compt.ones', 'import arrayblow as ab\n'), (782, 'arrayblow.v1.compt.ones', 'ab.v1.compt.ones', 'import arrayblow as ab\n'), (725, 'arrayblow.v1.compt.abs', 'ab.v1.compt.abs', 'import arrayblow as ab\n')]