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).
docstr: Cast the floating-point `parmas` to `jax.numpy.float16`. This returns a new `params` tree and does not cast the
  `params` in place.
  
  This method can be used on GPU to explicitly convert the model parameters to float16 precision to do full
  half-precision training or to save weights in float16 for inference in order to save memory and improve speed.
  
  Arguments:
      params (`Union[Dict, FrozenDict]`):
          A `PyTree` of model parameters.
      mask (`Union[Dict, FrozenDict]`):
          A `PyTree` with same structure as the `params` tree. The leaves should be booleans, `True` for params
          you want to cast, and should be `False` for those you want to skip
  
  Examples:
  
  ```python
  >>> from transformers import FlaxBertModel
  
  >>> # load model
  >>> model = FlaxBertModel.from_pretrained("bert-base-cased")
  >>> # By default, the model params will be in fp32, to cast these to float16
  >>> model.params = model.to_fp16(model.params)
  >>> # If you want don't want to cast certain parameters (for example layer norm bias and scale)
  >>> # then pass the mask as follows
  >>> from flax import traverse_util
  
  >>> model = FlaxBertModel.from_pretrained("bert-base-cased")
  >>> flat_params = traverse_util.flatten_dict(model.params)
  >>> mask = {
  ...     path: (path[-2] != ("LayerNorm", "bias") and path[-2:] != ("LayerNorm", "scale"))
  ...     for path in flat_params
  ... }
  >>> mask = traverse_util.unflatten_dict(mask)
  >>> model.params = model.to_fp16(model.params, mask)
  ```
docstr_tok: ["Cast","the","floating-point","`","parmas","`","to","`","jax.numpy.float16","`",".","This","returns","a","new","`","params","`","tree","and","does","not","cast","the","`","params","`","in","place",".","This","method","can","be","used","on","GPU","to","explicitly","convert","the","model","parameters","to","float16","precision","to","do","full","half-precision","training","or","to","save","weights","in","float16","for","inference","in","order","to","save","memory","and","improve","speed",".","Arguments",":","params","(","`","Union","[","Dict",",","FrozenDict","]","`",")",":","A","`","PyTree","`","of","model","parameters",".","mask","(","`","Union","[","Dict",",","FrozenDict","]","`",")",":","A","`","PyTree","`","with","same","structure","as","the","`","params","`","tree",".","The","leaves","should","be","booleans",",","`","True","`","for","params","you","want","to","cast",",","and","should","be","`","False","`","for","those","you","want","to","skip","Examples",":","``","`","python",">",">",">","from","transformers","import","FlaxBertModel",">",">",">","#","load","model",">",">",">","model","=","FlaxBertModel.from_pretrained","(","``","bert-base-cased","''",")",">",">",">","#","By","default",",","the","model","params","will","be","in","fp32",",","to","cast","these","to","float16",">",">",">","model.params","=","model.to_fp16","(","model.params",")",">",">",">","#","If","you","want","do","n't","want","to","cast","certain","parameters","(","for","example","layer","norm","bias","and","scale",")",">",">",">","#","then","pass","the","mask","as","follows",">",">",">","from","flax","import","traverse_util",">",">",">","model","=","FlaxBertModel.from_pretrained","(","``","bert-base-cased","''",")",">",">",">","flat_params","=","traverse_util.flatten_dict","(","model.params",")",">",">",">","mask","=","{","...","path",":","(","path","[","-2","]","!","=","(","``","LayerNorm","''",",","``","bias","''",")","and","path","[","-2",":","]","!","=","(","``","LayerNorm","''",",","``","scale","''",")",")","...","for","path","in","flat_params","...","}",">",">",">","mask","=","traverse_util.unflatten_dict","(","mask",")",">",">",">","model.params","=","model.to_fp16","(","model.params",",","mask",")","``","`"]
code: def to_fp16(self, params: Union[Dict, FrozenDict], mask: Any = None):
      r"""
      Cast the floating-point `parmas` to `jax.numpy.float16`. This returns a new `params` tree and does not cast the
      `params` in place.
  
      This method can be used on GPU to explicitly convert the model parameters to float16 precision to do full
      half-precision training or to save weights in float16 for inference in order to save memory and improve speed.
  
      Arguments:
          params (`Union[Dict, FrozenDict]`):
              A `PyTree` of model parameters.
          mask (`Union[Dict, FrozenDict]`):
              A `PyTree` with same structure as the `params` tree. The leaves should be booleans, `True` for params
              you want to cast, and should be `False` for those you want to skip
  
      Examples:
  
      ```python
      >>> from transformers import FlaxBertModel
  
      >>> # load model
      >>> model = FlaxBertModel.from_pretrained("bert-base-cased")
      >>> # By default, the model params will be in fp32, to cast these to float16
      >>> model.params = model.to_fp16(model.params)
      >>> # If you want don't want to cast certain parameters (for example layer norm bias and scale)
      >>> # then pass the mask as follows
      >>> from flax import traverse_util
  
      >>> model = FlaxBertModel.from_pretrained("bert-base-cased")
      >>> flat_params = traverse_util.flatten_dict(model.params)
      >>> mask = {
      ...     path: (path[-2] != ("LayerNorm", "bias") and path[-2:] != ("LayerNorm", "scale"))
      ...     for path in flat_params
      ... }
      >>> mask = traverse_util.unflatten_dict(mask)
      >>> model.params = model.to_fp16(model.params, mask)
      ```"""
      return self._cast_floating_to(params, jnp.float16, mask)
  
imports: import gc
  import json
  import os
  import re
  import warnings
  from functools import partial
  from pickle import UnpicklingError
  from typing import Any, Dict, Optional, 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 import FlaxGenerationMixin, GenerationConfig
  from .modeling_flax_pytorch_utils import load_pytorch_checkpoint_in_flax_state_dict
  from .utils import FLAX_WEIGHTS_INDEX_NAME, FLAX_WEIGHTS_NAME, SAFE_WEIGHTS_INDEX_NAME, SAFE_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
  from .utils.import_utils import is_safetensors_available
label_desc: Use image node_id 15 for calling the FlaxPreTrainedModel obj's underlying member method code with example usage: obj.to_fp16(params, mask) and returns: self