
"""simple docstring"""







import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_mpnet import MPNetTokenizer


lowercase_						    =					logging.get_logger(__name__)

lowercase_						    =					{"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}

lowercase_						    =					{
    "vocab_file": {
        "microsoft/mpnet-base": "https://huggingface.co/microsoft/mpnet-base/resolve/main/vocab.txt",
    },
    "tokenizer_file": {
        "microsoft/mpnet-base": "https://huggingface.co/microsoft/mpnet-base/resolve/main/tokenizer.json",
    },
}

lowercase_						    =					{
    "microsoft/mpnet-base": 512,
}

lowercase_						    =					{
    "microsoft/mpnet-base": {"do_lower_case": True},
}



class 		UpperCAmelCase_	(lowerCamelCase_						):
 """simple docstring"""




 UpperCamelCase_						:     List[Any]				   =    VOCAB_FILES_NAMES
 UpperCamelCase_						:     Union[str, Any]				   =    PRETRAINED_VOCAB_FILES_MAP
 UpperCamelCase_						:     Union[str, Any]				   =    PRETRAINED_INIT_CONFIGURATION
 UpperCamelCase_						:     str				   =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 UpperCamelCase_						:     Union[str, Any]				   =    MPNetTokenizer
 UpperCamelCase_						:     Optional[Any]				   =    ["""input_ids""", """attention_mask"""]




 def __init__( self       : Tuple     ,   a_       : Union[str, Any]=None     ,   a_       : Union[str, Any]=None     ,   a_       : Optional[int]=True     ,   a_       : str="<s>"     ,   a_       : int="</s>"     ,   a_       : Union[str, Any]="</s>"     ,   a_       : List[str]="<s>"     ,   a_       : List[str]="[UNK]"     ,   a_       : str="<pad>"     ,   a_       : Optional[int]="<mask>"     ,   a_       : Dict=True     ,   a_       : Union[str, Any]=None     ,   **a_       : Optional[Any]     ,   )->					Optional[Any]:



      """simple docstring"""

      super().__init__(
          a_     ,   tokenizer_file=a_     ,   do_lower_case=a_     ,   bos_token=a_     ,   eos_token=a_     ,   sep_token=a_     ,   cls_token=a_     ,   unk_token=a_     ,   pad_token=a_     ,   mask_token=a_     ,   tokenize_chinese_chars=a_     ,   strip_accents=a_     ,   **a_     ,   )

      UpperCAmelCase_						:    Union[str, Any]    					=  json.loads(self.backend_tokenizer.normalizer.__getstate__() )
      if (
          pre_tok_state.get("""lowercase"""     ,   a_ ) != do_lower_case
          or pre_tok_state.get("""strip_accents"""     ,   a_ ) != strip_accents
      ):
           UpperCAmelCase_						:    List[Any]    					=  getattr(a_     ,   pre_tok_state.pop("""type""" ) )
           UpperCAmelCase_						:    Optional[int]    					=  do_lower_case
           UpperCAmelCase_						:    Any    					=  strip_accents
           UpperCAmelCase_						:    List[Any]    					=  pre_tok_class(**a_ )

      UpperCAmelCase_						:    Optional[int]    					=  do_lower_case




 @property
 def   a      ( self       : List[Any] )->					str:



      """simple docstring"""

      if self._mask_token is None:
           if self.verbose:
                logger.error("""Using mask_token, but it is not set yet.""" )
           return None
      return str(self._mask_token )




 @mask_token.setter
 def   a      ( self       : Any     ,   a_       : Optional[int] )->					str:



      """simple docstring"""

      UpperCAmelCase_						:    Dict    					=  AddedToken(a_     ,   lstrip=a_     ,   rstrip=a_ ) if isinstance(a_     ,   a_ ) else value
      UpperCAmelCase_						:    Optional[Any]    					=  value




 def   a      ( self       : List[str]     ,   a_       : Dict     ,   a_       : List[str]=None )->					Optional[int]:



      """simple docstring"""

      UpperCAmelCase_						:    List[str]    					=  [self.bos_token_id] + token_ids_a + [self.eos_token_id]
      if token_ids_a is None:
           return output

      return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]




 def   a      ( self       : List[str]     ,   a_       : List[int]     ,   a_       : Optional[List[int]] = None )->					List[int]:



      """simple docstring"""

      UpperCAmelCase_						:    List[str]    					=  [self.sep_token_id]
      UpperCAmelCase_						:    Dict    					=  [self.cls_token_id]

      if token_ids_a is None:
           return len(cls + token_ids_a + sep ) * [0]
      return len(cls + token_ids_a + sep + sep + token_ids_a + sep ) * [0]






 def   a      ( self       : Union[str, Any]     ,   a_       : str     ,   a_       : Optional[str] = None )->					Tuple[str]:



      """simple docstring"""

      UpperCAmelCase_						:    Dict    					=  self._tokenizer.model.save(a_     ,   name=a_ )
      return tuple(a_ )



"""simple docstring"""







import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import PriorTransformer, UnCLIPPipeline, UnCLIPScheduler, UNetaDConditionModel, UNetaDModel
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()



class 		UpperCAmelCase_	(lowerCamelCase_  , unittest.TestCase						):
 """simple docstring"""




 UpperCamelCase_						:     List[str]				   =    UnCLIPPipeline
 UpperCamelCase_						:     Any				   =    TEXT_TO_IMAGE_PARAMS - {
     """negative_prompt""",
     """height""",
     """width""",
     """negative_prompt_embeds""",
     """guidance_scale""",
     """prompt_embeds""",
     """cross_attention_kwargs""",
 }
 UpperCamelCase_						:     int				   =    TEXT_TO_IMAGE_BATCH_PARAMS
 UpperCamelCase_						:     int				   =    [
     """generator""",
     """return_dict""",
     """prior_num_inference_steps""",
     """decoder_num_inference_steps""",
     """super_res_num_inference_steps""",
 ]
 UpperCamelCase_						:     Union[str, Any]				   =    False




 @property
 def   a      ( self       : Tuple )->					str:



      """simple docstring"""

      return 32




 @property
 def   a      ( self       : Tuple )->					Any:



      """simple docstring"""

      return 32




 @property
 def   a      ( self       : Union[str, Any] )->					str:



      """simple docstring"""

      return self.time_input_dim




 @property
 def   a      ( self       : int )->					Dict:



      """simple docstring"""

      return self.time_input_dim * 4




 @property
 def   a      ( self       : List[str] )->					Tuple:



      """simple docstring"""

      return 1_00




 @property
 def   a      ( self       : Tuple )->					Optional[Any]:



      """simple docstring"""

      UpperCAmelCase_						:    List[Any]    					=  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""" )
      return tokenizer




 @property
 def   a      ( self       : Optional[Any] )->					Dict:



      """simple docstring"""

      torch.manual_seed(0 )
      UpperCAmelCase_						:    Any    					=  CLIPTextConfig(
          bos_token_id=0     ,   eos_token_id=2     ,   hidden_size=self.text_embedder_hidden_size     ,   projection_dim=self.text_embedder_hidden_size     ,   intermediate_size=37     ,   layer_norm_eps=1E-05     ,   num_attention_heads=4     ,   num_hidden_layers=5     ,   pad_token_id=1     ,   vocab_size=10_00     ,   )
      return CLIPTextModelWithProjection(a_ )




 @property
 def   a      ( self       : str )->					Optional[Any]:



      """simple docstring"""

      torch.manual_seed(0 )

      UpperCAmelCase_						:    Optional[Any]    					=  {
          """num_attention_heads""": 2,
          """attention_head_dim""": 12,
          """embedding_dim""": self.text_embedder_hidden_size,
          """num_layers""": 1,
      }

      UpperCAmelCase_						:    Tuple    					=  PriorTransformer(**a_ )
      return model




 @property
 def   a      ( self       : List[str] )->					Any:



      """simple docstring"""

      torch.manual_seed(0 )

      UpperCAmelCase_						:    Optional[int]    					=  {
          """clip_embeddings_dim""": self.text_embedder_hidden_size,
          """time_embed_dim""": self.time_embed_dim,
          """cross_attention_dim""": self.cross_attention_dim,
      }

      UpperCAmelCase_						:    Tuple    					=  UnCLIPTextProjModel(**a_ )
      return model




 @property
 def   a      ( self       : Optional[int] )->					Tuple:



      """simple docstring"""

      torch.manual_seed(0 )

      UpperCAmelCase_						:    Tuple    					=  {
          """sample_size""": 32,
          # RGB in channels
          """in_channels""": 3,
          # Out channels is double in channels because predicts mean and variance
          """out_channels""": 6,
          """down_block_types""": ("""ResnetDownsampleBlock2D""", """SimpleCrossAttnDownBlock2D"""),
          """up_block_types""": ("""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""),
          """mid_block_type""": """UNetMidBlock2DSimpleCrossAttn""",
          """block_out_channels""": (self.block_out_channels_a, self.block_out_channels_a * 2),
          """layers_per_block""": 1,
          """cross_attention_dim""": self.cross_attention_dim,
          """attention_head_dim""": 4,
          """resnet_time_scale_shift""": """scale_shift""",
          """class_embed_type""": """identity""",
      }

      UpperCAmelCase_						:    List[str]    					=  UNetaDConditionModel(**a_ )
      return model




 @property
 def   a      ( self       : Dict )->					Optional[int]:



      """simple docstring"""

      return {
          "sample_size": 64,
          "layers_per_block": 1,
          "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
          "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
          "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
          "in_channels": 6,
          "out_channels": 3,
      }




 @property
 def   a      ( self       : str )->					Optional[int]:



      """simple docstring"""

      torch.manual_seed(0 )

      UpperCAmelCase_						:    Tuple    					=  UNetaDModel(**self.dummy_super_res_kwargs )
      return model




 @property
 def   a      ( self       : Union[str, Any] )->					Union[str, Any]:



      """simple docstring"""

      torch.manual_seed(1 )

      UpperCAmelCase_						:    Dict    					=  UNetaDModel(**self.dummy_super_res_kwargs )
      return model




 def   a      ( self       : Tuple )->					str:



      """simple docstring"""

      UpperCAmelCase_						:    List[str]    					=  self.dummy_prior
      UpperCAmelCase_						:    Any    					=  self.dummy_decoder
      UpperCAmelCase_						:    Union[str, Any]    					=  self.dummy_text_proj
      UpperCAmelCase_						:    int    					=  self.dummy_text_encoder
      UpperCAmelCase_						:    Union[str, Any]    					=  self.dummy_tokenizer
      UpperCAmelCase_						:    Tuple    					=  self.dummy_super_res_first
      UpperCAmelCase_						:    List[str]    					=  self.dummy_super_res_last

      UpperCAmelCase_						:    Dict    					=  UnCLIPScheduler(
          variance_type="""fixed_small_log"""     ,   prediction_type="""sample"""     ,   num_train_timesteps=10_00     ,   clip_sample_range=5.0     ,   )

      UpperCAmelCase_						:    List[str]    					=  UnCLIPScheduler(
          variance_type="""learned_range"""     ,   prediction_type="""epsilon"""     ,   num_train_timesteps=10_00     ,   )

      UpperCAmelCase_						:    List[str]    					=  UnCLIPScheduler(
          variance_type="""fixed_small_log"""     ,   prediction_type="""epsilon"""     ,   num_train_timesteps=10_00     ,   )

      UpperCAmelCase_						:    List[str]    					=  {
          """prior""": prior,
          """decoder""": decoder,
          """text_proj""": text_proj,
          """text_encoder""": text_encoder,
          """tokenizer""": tokenizer,
          """super_res_first""": super_res_first,
          """super_res_last""": super_res_last,
          """prior_scheduler""": prior_scheduler,
          """decoder_scheduler""": decoder_scheduler,
          """super_res_scheduler""": super_res_scheduler,
      }

      return components




 def   a      ( self       : Union[str, Any]     ,   a_       : Optional[Any]     ,   a_       : Dict=0 )->					Union[str, Any]:



      """simple docstring"""

      if str(a_ ).startswith("""mps""" ):
           UpperCAmelCase_						:    Union[str, Any]    					=  torch.manual_seed(a_ )
      else:
           UpperCAmelCase_						:    List[str]    					=  torch.Generator(device=a_ ).manual_seed(a_ )
      UpperCAmelCase_						:    Dict    					=  {
          """prompt""": """horse""",
          """generator""": generator,
          """prior_num_inference_steps""": 2,
          """decoder_num_inference_steps""": 2,
          """super_res_num_inference_steps""": 2,
          """output_type""": """numpy""",
      }
      return inputs




 def   a      ( self       : str )->					Optional[int]:



      """simple docstring"""

      UpperCAmelCase_						:    List[Any]    					=  """cpu"""

      UpperCAmelCase_						:    int    					=  self.get_dummy_components()

      UpperCAmelCase_						:    str    					=  self.pipeline_class(**a_ )
      UpperCAmelCase_						:    Any    					=  pipe.to(a_ )

      pipe.set_progress_bar_config(disable=a_ )

      UpperCAmelCase_						:    int    					=  pipe(**self.get_dummy_inputs(a_ ) )
      UpperCAmelCase_						:    List[str]    					=  output.images

      UpperCAmelCase_						:    Optional[int]    					=  pipe(
          **self.get_dummy_inputs(a_ )     ,   return_dict=a_     ,   )[0]

      UpperCAmelCase_						:    List[str]    					=  image[0, -3:, -3:, -1]
      UpperCAmelCase_						:    int    					=  image_from_tuple[0, -3:, -3:, -1]

      assert image.shape == (1, 64, 64, 3)

      UpperCAmelCase_						:    Any    					=  np.array(
          [
              0.9997,
              0.9988,
              0.0028,
              0.9997,
              0.9984,
              0.9965,
              0.0029,
              0.9986,
              0.0025,
          ] )

      assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2
      assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() < 1E-2




 def   a      ( self       : int )->					str:



      """simple docstring"""

      UpperCAmelCase_						:    str    					=  torch.device("""cpu""" )



      class 		UpperCAmelCase_	:
       """simple docstring"""




       UpperCamelCase_						:     Optional[Any]				   =    1

      UpperCAmelCase_						:    List[str]    					=  self.get_dummy_components()

      UpperCAmelCase_						:    str    					=  self.pipeline_class(**a_ )
      UpperCAmelCase_						:    Any    					=  pipe.to(a_ )

      UpperCAmelCase_						:    List[Any]    					=  components["""prior"""]
      UpperCAmelCase_						:    str    					=  components["""decoder"""]
      UpperCAmelCase_						:    str    					=  components["""super_res_first"""]
      UpperCAmelCase_						:    Optional[int]    					=  components["""tokenizer"""]
      UpperCAmelCase_						:    List[str]    					=  components["""text_encoder"""]

      UpperCAmelCase_						:    Any    					=  torch.Generator(device=a_ ).manual_seed(0 )
      UpperCAmelCase_						:    Optional[Any]    					=  prior.dtype
      UpperCAmelCase_						:    List[Any]    					=  1

      UpperCAmelCase_						:    int    					=  (batch_size, prior.config.embedding_dim)
      UpperCAmelCase_						:    Optional[Any]    					=  pipe.prepare_latents(
          a_     ,   dtype=a_     ,   device=a_     ,   generator=a_     ,   latents=a_     ,   scheduler=DummyScheduler() )
      UpperCAmelCase_						:    List[str]    					=  (batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
      UpperCAmelCase_						:    int    					=  pipe.prepare_latents(
          a_     ,   dtype=a_     ,   device=a_     ,   generator=a_     ,   latents=a_     ,   scheduler=DummyScheduler() )

      UpperCAmelCase_						:    List[str]    					=  (
          batch_size,
          super_res_first.config.in_channels // 2,
          super_res_first.config.sample_size,
          super_res_first.config.sample_size,
      )
      UpperCAmelCase_						:    int    					=  pipe.prepare_latents(
          a_     ,   dtype=a_     ,   device=a_     ,   generator=a_     ,   latents=a_     ,   scheduler=DummyScheduler() )

      pipe.set_progress_bar_config(disable=a_ )

      UpperCAmelCase_						:    str    					=  """this is a prompt example"""

      UpperCAmelCase_						:    Any    					=  torch.Generator(device=a_ ).manual_seed(0 )
      UpperCAmelCase_						:    int    					=  pipe(
          [prompt]     ,   generator=a_     ,   prior_num_inference_steps=2     ,   decoder_num_inference_steps=2     ,   super_res_num_inference_steps=2     ,   prior_latents=a_     ,   decoder_latents=a_     ,   super_res_latents=a_     ,   output_type="""np"""     ,   )
      UpperCAmelCase_						:    Tuple    					=  output.images

      UpperCAmelCase_						:    int    					=  tokenizer(
          a_     ,   padding="""max_length"""     ,   max_length=tokenizer.model_max_length     ,   return_tensors="""pt"""     ,   )
      UpperCAmelCase_						:    Union[str, Any]    					=  text_encoder(text_inputs.input_ids )
      UpperCAmelCase_						:    Optional[Any]    					=  text_inputs.attention_mask

      UpperCAmelCase_						:    List[str]    					=  torch.Generator(device=a_ ).manual_seed(0 )
      UpperCAmelCase_						:    List[str]    					=  pipe(
          generator=a_     ,   prior_num_inference_steps=2     ,   decoder_num_inference_steps=2     ,   super_res_num_inference_steps=2     ,   prior_latents=a_     ,   decoder_latents=a_     ,   super_res_latents=a_     ,   text_model_output=a_     ,   text_attention_mask=a_     ,   output_type="""np"""     ,   )[0]

      # make sure passing text embeddings manually is identical
      assert np.abs(image - image_from_text ).max() < 1E-4




 @skip_mps
 def   a      ( self       : Tuple )->					Tuple:



      """simple docstring"""

      UpperCAmelCase_						:    int    					=  torch_device == """cpu"""

      self._test_attention_slicing_forward_pass(test_max_difference=a_     ,   expected_max_diff=0.01 )




 @skip_mps
 def   a      ( self       : Dict )->					Union[str, Any]:



      """simple docstring"""

      UpperCAmelCase_						:    Optional[Any]    					=  torch_device == """cpu"""
      UpperCAmelCase_						:    List[str]    					=  True
      UpperCAmelCase_						:    Any    					=  [
          """prior_num_inference_steps""",
          """decoder_num_inference_steps""",
          """super_res_num_inference_steps""",
      ]

      self._test_inference_batch_single_identical(
          test_max_difference=a_     ,   relax_max_difference=a_     ,   additional_params_copy_to_batched_inputs=a_     ,   )




 def   a      ( self       : Any )->					Optional[Any]:



      """simple docstring"""

      UpperCAmelCase_						:    Union[str, Any]    					=  [
          """prior_num_inference_steps""",
          """decoder_num_inference_steps""",
          """super_res_num_inference_steps""",
      ]

      if torch_device == "mps":
           # TODO: MPS errors with larger batch sizes
           UpperCAmelCase_						:    List[Any]    					=  [2, 3]
           self._test_inference_batch_consistent(
               batch_sizes=a_     ,   additional_params_copy_to_batched_inputs=a_     ,   )
      else:
           self._test_inference_batch_consistent(
               additional_params_copy_to_batched_inputs=a_ )




 @skip_mps
 def   a      ( self       : Optional[int] )->					Any:



      """simple docstring"""

      return super().test_dict_tuple_outputs_equivalent()




 @skip_mps
 def   a      ( self       : int )->					Any:



      """simple docstring"""

      return super().test_save_load_local()






 @skip_mps
 def   a      ( self       : List[Any] )->					List[Any]:



      """simple docstring"""

      return super().test_save_load_optional_components()



@nightly
class 		UpperCAmelCase_	(unittest.TestCase						):
 """simple docstring"""




 def   a      ( self       : Optional[int] )->					Union[str, Any]:



      """simple docstring"""

      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()






 def   a      ( self       : int )->					Tuple:



      """simple docstring"""

      UpperCAmelCase_						:    List[Any]    					=  load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
          """/unclip/karlo_v1_alpha_horse_cpu.npy""" )

      UpperCAmelCase_						:    List[Any]    					=  UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha""" )
      pipeline.set_progress_bar_config(disable=a_ )

      UpperCAmelCase_						:    List[str]    					=  torch.manual_seed(0 )
      UpperCAmelCase_						:    Any    					=  pipeline(
          """horse"""     ,   num_images_per_prompt=1     ,   generator=a_     ,   output_type="""np"""     ,   )

      UpperCAmelCase_						:    Optional[int]    					=  output.images[0]

      assert image.shape == (2_56, 2_56, 3)
      assert np.abs(expected_image - image ).max() < 1E-1



@slow
@require_torch_gpu
class 		UpperCAmelCase_	(unittest.TestCase						):
 """simple docstring"""




 def   a      ( self       : str )->					Any:



      """simple docstring"""

      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()




 def   a      ( self       : List[str] )->					List[Any]:



      """simple docstring"""

      UpperCAmelCase_						:    Tuple    					=  load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
          """/unclip/karlo_v1_alpha_horse_fp16.npy""" )

      UpperCAmelCase_						:    Optional[int]    					=  UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""     ,   torch_dtype=torch.floataa )
      UpperCAmelCase_						:    Optional[int]    					=  pipeline.to(a_ )
      pipeline.set_progress_bar_config(disable=a_ )

      UpperCAmelCase_						:    Any    					=  torch.Generator(device="""cpu""" ).manual_seed(0 )
      UpperCAmelCase_						:    Union[str, Any]    					=  pipeline(
          """horse"""     ,   generator=a_     ,   output_type="""np"""     ,   )

      UpperCAmelCase_						:    Optional[Any]    					=  output.images[0]

      assert image.shape == (2_56, 2_56, 3)

      assert_mean_pixel_difference(a_     ,   a_ )






 def   a      ( self       : Optional[int] )->					int:



      """simple docstring"""

      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()

      UpperCAmelCase_						:    Optional[Any]    					=  UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""     ,   torch_dtype=torch.floataa )
      UpperCAmelCase_						:    List[str]    					=  pipe.to(a_ )
      pipe.set_progress_bar_config(disable=a_ )
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      UpperCAmelCase_						:    int    					=  pipe(
          """horse"""     ,   num_images_per_prompt=1     ,   prior_num_inference_steps=2     ,   decoder_num_inference_steps=2     ,   super_res_num_inference_steps=2     ,   output_type="""np"""     ,   )

      UpperCAmelCase_						:    Optional[Any]    					=  torch.cuda.max_memory_allocated()
      # make sure that less than 7 GB is allocated
      assert mem_bytes < 7 * 10**9

