class_docstr: Base class for all models.
  
  [`FlaxPreTrainedModel`] takes care of storing the configuration of the models and handles methods for loading,
  downloading and saving models.
  
  Class attributes (overridden by derived classes):
  
      - **config_class** ([`PretrainedConfig`]) -- A subclass of [`PretrainedConfig`] to use as configuration class
        for this model architecture.
      - **base_model_prefix** (`str`) -- A string indicating the attribute associated to the base model in derived
        classes of the same architecture adding modules on top of the base model.
      - **main_input_name** (`str`) -- The name of the principal input to the model (often `input_ids` for NLP
        models, `pixel_values` for vision models and `input_values` for speech models).
class_docstr_tok: ["Base","class","for","all","models",".","[","`","FlaxPreTrainedModel","`","]","takes","care","of","storing","the","configuration","of","the","models","and","handles","methods","for","loading",",","downloading","and","saving","models",".","Class","attributes","(","overridden","by","derived","classes",")",":","-","*","*","config_class","*","*","(","[","`","PretrainedConfig","`","]",")","--","A","subclass","of","[","`","PretrainedConfig","`","]","to","use","as","configuration","class","for","this","model","architecture",".","-","*","*","base_model_prefix","*","*","(","`","str","`",")","--","A","string","indicating","the","attribute","associated","to","the","base","model","in","derived","classes","of","the","same","architecture","adding","modules","on","top","of","the","base","model",".","-","*","*","main_input_name","*","*","(","`","str","`",")","--","The","name","of","the","principal","input","to","the","model","(","often","`","input_ids","`","for","NLP","models",",","`","pixel_values","`","for","vision","models","and","`","input_values","`","for","speech","models",")","."]
docstr: Save a model and its configuration file to a directory, so that it can be re-loaded using the
  `[`~FlaxPreTrainedModel.from_pretrained`]` class method
  
  Arguments:
      save_directory (`str` or `os.PathLike`):
          Directory to which to save. Will be created if it doesn't exist.
      push_to_hub (`bool`, *optional*, defaults to `False`):
          Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
          repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
          namespace).
      max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`):
          The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size
          lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5MB"`).
  
          <Tip warning={true}>
  
          If a single weight of the model is bigger than `max_shard_size`, it will be in its own checkpoint shard
          which will be bigger than `max_shard_size`.
  
          </Tip>
  
      kwargs:
          Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
docstr_tok: ["Save","a","model","and","its","configuration","file","to","a","directory",",","so","that","it","can","be","re-loaded","using","the","`","[","`","~FlaxPreTrainedModel.from_pretrained","`","]","`","class","method","Arguments",":","save_directory","(","`","str","`","or","`","os.PathLike","`",")",":","Directory","to","which","to","save",".","Will","be","created","if","it","does","n't","exist",".","push_to_hub","(","`","bool","`",",","*","optional","*",",","defaults","to","`","False","`",")",":","Whether","or","not","to","push","your","model","to","the","Hugging","Face","model","hub","after","saving","it",".","You","can","specify","the","repository","you","want","to","push","to","with","`","repo_id","`","(","will","default","to","the","name","of","`","save_directory","`","in","your","namespace",")",".","max_shard_size","(","`","int","`","or","`","str","`",",","*","optional","*",",","defaults","to","`","``","10GB","''","`",")",":","The","maximum","size","for","a","checkpoint","before","being","sharded",".","Checkpoints","shard","will","then","be","each","of","size","lower","than","this","size",".","If","expressed","as","a","string",",","needs","to","be","digits","followed","by","a","unit","(","like","`","``","5MB","''","`",")",".","<","Tip","warning=","{","true","}",">","If","a","single","weight","of","the","model","is","bigger","than","`","max_shard_size","`",",","it","will","be","in","its","own","checkpoint","shard","which","will","be","bigger","than","`","max_shard_size","`",".","<","\/Tip",">","kwargs",":","Additional","key","word","arguments","passed","along","to","the","[","`","~utils.PushToHubMixin.push_to_hub","`","]","method","."]
code: def save_pretrained(
      self,
      save_directory: Union[str, os.PathLike],
      params=None,
      push_to_hub=False,
      max_shard_size="10GB",
      **kwargs,
  ):
      """
      Save a model and its configuration file to a directory, so that it can be re-loaded using the
      `[`~FlaxPreTrainedModel.from_pretrained`]` class method
  
      Arguments:
          save_directory (`str` or `os.PathLike`):
              Directory to which to save. Will be created if it doesn't exist.
          push_to_hub (`bool`, *optional*, defaults to `False`):
              Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
              repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
              namespace).
          max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`):
              The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size
              lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5MB"`).
  
              <Tip warning={true}>
  
              If a single weight of the model is bigger than `max_shard_size`, it will be in its own checkpoint shard
              which will be bigger than `max_shard_size`.
  
              </Tip>
  
          kwargs:
              Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
      """
      if os.path.isfile(save_directory):
          logger.error(
              f"Provided path ({save_directory}) should be a directory, not a file"
          )
          return
  
      os.makedirs(save_directory, exist_ok=True)
  
      if push_to_hub:
          commit_message = kwargs.pop("commit_message", None)
          repo_id = kwargs.pop(
              "repo_id", save_directory.split(os.path.sep)[-1]
          )
          repo_id, token = self._create_repo(repo_id, **kwargs)
          files_timestamps = self._get_files_timestamps(save_directory)
  
      # get abs dir
      save_directory = os.path.abspath(save_directory)
      # save config as well
      self.config.architectures = [self.__class__.__name__[4:]]
  
      # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be
      # loaded from the Hub.
      if self._auto_class is not None:
          custom_object_save(self, save_directory, config=self.config)
  
      self.config.save_pretrained(save_directory)
  
      # save model
      output_model_file = os.path.join(
          save_directory, FLAX_WEIGHTS_NAME
      )
  
      shards, index = flax_shard_checkpoint(
          params if params is not None else self.params, max_shard_size
      )
      # Clean the folder from a previous save
      for filename in os.listdir(save_directory):
          full_filename = os.path.join(save_directory, filename)
          if (
              filename.startswith(FLAX_WEIGHTS_NAME[:-4])
              and os.path.isfile(full_filename)
              and filename not in shards.keys()
          ):
              os.remove(full_filename)
  
      if index is None:
          with open(output_model_file, "wb") as f:
              params = params if params is not None else self.params
              model_bytes = to_bytes(params)
              f.write(model_bytes)
  
      else:
          save_index_file = os.path.join(
              save_directory, FLAX_WEIGHTS_INDEX_NAME
          )
          # Save the index as well
          with open(save_index_file, "w", encoding="utf-8") as f:
              content = (
                  json.dumps(index, indent=2, sort_keys=True) + "\n"
              )
              f.write(content)
          logger.info(
              f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
              f"split in {len(shards)} checkpoint shards. You can find where each parameters has been saved in the "
              f"index located at {save_index_file}."
          )
          for shard_file, shard in shards.items():
              # the shard item are unflattened, to save them we need to flatten them again
              with open(
                  os.path.join(save_directory, shard_file), mode="wb"
              ) as f:
                  params = unflatten_dict(shard, sep="/")
                  shard_bytes = to_bytes(params)
                  f.write(shard_bytes)
  
      logger.info(f"Model weights saved in {output_model_file}")
  
      if push_to_hub:
          self._upload_modified_files(
              save_directory,
              repo_id,
              files_timestamps,
              commit_message=commit_message,
              token=token,
          )
  
imports: import gc
  import json
  import os
  import re
  from functools import partial
  from pickle import UnpicklingError
  from typing import Any, Dict, Set, Tuple, Union
  import flax.linen
  import jax
  import jax.numpy
  import msgpack.exceptions
  from flax.core.frozen_dict import FrozenDict, unfreeze
  from flax.serialization import from_bytes, to_bytes
  from flax.traverse_util import flatten_dict, unflatten_dict
  from jax.random import PRNGKey
  from .configuration_utils import PretrainedConfig
  from .dynamic_module_utils import custom_object_save
  from .generation_flax_utils import FlaxGenerationMixin
  from .modeling_flax_pytorch_utils import load_pytorch_checkpoint_in_flax_state_dict
  from .utils import FLAX_WEIGHTS_INDEX_NAME, FLAX_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, WEIGHTS_NAME, PushToHubMixin, add_code_sample_docstrings, add_start_docstrings_to_model_forward, cached_file, copy_func, download_url, has_file, is_offline_mode, is_remote_url, logging, replace_return_docstrings
  from .utils.hub import convert_file_size_to_int, get_checkpoint_shard_files
label_desc: Use image node_id 18 for calling the FlaxPreTrainedModel obj's underlying member method code with example usage: obj.save_pretrained(save_directory, params, push_to_hub, max_shard_size) without return types