code: def test_gradient_checkpointing(self):
      (
          config,
          inputs_dict,
      ) = self.model_tester.prepare_config_and_inputs_for_common()
  
      for model_class in self.all_model_classes:
          # prepare inputs
          prepared_inputs_dict = self._prepare_for_class(
              inputs_dict, model_class
          )
          model = model_class(config)
          remat_model = model_class(config)
  
          try:
              remat_model.enable_gradient_checkpointing()
          except NotImplementedError:
              continue
  
          outputs = model(**prepared_inputs_dict)
          remat_outputs = remat_model(**prepared_inputs_dict)
  
          # ensure that the dicts of outputs contain the same keys
          self.assertEqual(outputs.keys(), remat_outputs.keys())
  
          outputs = outputs.to_tuple()
          remat_outputs = remat_outputs.to_tuple()
  
          # ensure that the outputs remain precisely equal
          for output, remat_output in zip(outputs, remat_outputs):
              self.assertTrue((output == remat_output).all())
  
imports: import copy
  import inspect
  import json
  import random
  import tempfile
  import unittest
  from typing import List, Tuple
  import numpy
  import transformers
  from huggingface_hub import HfFolder, delete_repo, set_access_token
  from requests.exceptions import HTTPError
  from transformers import BertConfig, is_flax_available, is_torch_available
  from transformers.models.auto import get_values
  from transformers.testing_utils import TOKEN, USER, CaptureLogger, is_pt_flax_cross_test, is_staging_test, require_flax, torch_device
  from transformers.utils import logging
  from transformers.utils.generic import ModelOutput
label_desc: Use image node_id 32 for calling the FlaxModelTesterMixin obj's underlying member method code with example usage: obj.test_gradient_checkpointing() without return types