class ModelCheckpoint(Callback):

    def __init__(self, filepath, monitor='val_loss', verbose=0, save_best_only=False, save_weights_only=False, mode='auto', save_freq='epoch', options=None, **kwargs):
        super(ModelCheckpoint, self).__init__()
        self._supports_tf_logs = True
        self.monitor = monitor
        self.verbose = verbose
        self.filepath = path_to_string(filepath)
        self.save_best_only = save_best_only
        self.save_weights_only = save_weights_only
        self.save_freq = save_freq
        self.epochs_since_last_save = 0
        self._batches_seen_since_last_saving = 0
        self._last_batch_seen = 0
        if save_weights_only:
            if options is None or isinstance(options, checkpoint_options_lib.CheckpointOptions):
                self._options = options or checkpoint_options_lib.CheckpointOptions()
            else:
                raise TypeError('If save_weights_only is True, then `options` must be either None or a tf.train.CheckpointOptions')
        elif options is None or isinstance(options, save_options_lib.SaveOptions):
            self._options = options or save_options_lib.SaveOptions()
        else:
            raise TypeError('If save_weights_only is False, then `options` must beeither None or a tf.saved_model.SaveOptions')
        if 'load_weights_on_restart' in kwargs:
            self.load_weights_on_restart = kwargs['load_weights_on_restart']
            logging.warning('`load_weights_on_restart` argument is deprecated. Please use `model.load_weights()` for loading weights before the start of `model.fit()`.')
        else:
            self.load_weights_on_restart = False
        if 'period' in kwargs:
            self.period = kwargs['period']
            logging.warning('`period` argument is deprecated. Please use `save_freq` to specify the frequency in number of batches seen.')
        else:
            self.period = 1
        if mode not in ['auto', 'min', 'max']:
            logging.warning('ModelCheckpoint mode %s is unknown, fallback to auto mode.', mode)
            mode = 'auto'
        if mode == 'min':
            self.monitor_op = np.less
            self.best = np.inf
        elif mode == 'max':
            self.monitor_op = np.greater
            self.best = -np.inf
        elif 'acc' in self.monitor or self.monitor.startswith('fmeasure'):
            self.monitor_op = np.greater
            self.best = -np.inf
        else:
            self.monitor_op = np.less
            self.best = np.inf
        if self.save_freq != 'epoch' and (not isinstance(self.save_freq, int)):
            raise ValueError('Unrecognized save_freq: {}'.format(self.save_freq))
        self._chief_worker_only = False

    def on_train_begin(self, logs=None):
        if self.load_weights_on_restart:
            filepath_to_load = self._get_most_recently_modified_file_matching_pattern(self.filepath)
            if filepath_to_load is not None and self._checkpoint_exists(filepath_to_load):
                try:
                    self.model.load_weights(filepath_to_load)
                except (IOError, ValueError) as e:
                    raise ValueError('Error loading file from {}. Reason: {}'.format(filepath_to_load, e))

    def _implements_train_batch_hooks(self):
        return self.save_freq != 'epoch'

    def on_train_batch_end(self, batch, logs=None):
        if self._should_save_on_batch(batch):
            self._save_model(epoch=self._current_epoch, logs=logs)

    def on_epoch_begin(self, epoch, logs=None):
        self._current_epoch = epoch

    def on_epoch_end(self, epoch, logs=None):
        self.epochs_since_last_save += 1
        if self.save_freq == 'epoch':
            self._save_model(epoch=epoch, logs=logs)

    def _should_save_on_batch(self, batch):
        """Handles batch-level saving logic, supports steps_per_execution."""
        if self.save_freq == 'epoch':
            return False
        if batch <= self._last_batch_seen:
            add_batches = batch + 1
        else:
            add_batches = batch - self._last_batch_seen
        self._batches_seen_since_last_saving += add_batches
        self._last_batch_seen = batch
        if self._batches_seen_since_last_saving >= self.save_freq:
            self._batches_seen_since_last_saving = 0
            return True
        return False

    def _save_model(self, epoch, logs):
        """Saves the model.

    Args:
        epoch: the epoch this iteration is in.
        logs: the `logs` dict passed in to `on_batch_end` or `on_epoch_end`.
    """
        logs = logs or {}
        if isinstance(self.save_freq, int) or self.epochs_since_last_save >= self.period:
            logs = tf_utils.sync_to_numpy_or_python_type(logs)
            self.epochs_since_last_save = 0
            filepath = self._get_file_path(epoch, logs)
            try:
                if self.save_best_only:
                    current = logs.get(self.monitor)
                    if current is None:
                        logging.warning('Can save best model only with %s available, skipping.', self.monitor)
                    elif self.monitor_op(current, self.best):
                        if self.verbose > 0:
                            print('\nEpoch %05d: %s improved from %0.5f to %0.5f, saving model to %s' % (epoch + 1, self.monitor, self.best, current, filepath))
                        self.best = current
                        if self.save_weights_only:
                            self.model.save_weights(filepath, overwrite=True, options=self._options)
                        else:
                            self.model.save(filepath, overwrite=True, options=self._options)
                    elif self.verbose > 0:
                        print('\nEpoch %05d: %s did not improve from %0.5f' % (epoch + 1, self.monitor, self.best))
                else:
                    if self.verbose > 0:
                        print('\nEpoch %05d: saving model to %s' % (epoch + 1, filepath))
                    if self.save_weights_only:
                        self.model.save_weights(filepath, overwrite=True, options=self._options)
                    else:
                        self.model.save(filepath, overwrite=True, options=self._options)
                self._maybe_remove_file()
            except IsADirectoryError as e:
                raise IOError('Please specify a non-directory filepath for ModelCheckpoint. Filepath used is an existing directory: {}'.format(filepath))
            except IOError as e:
                if 'is a directory' in str(e.args[0]).lower():
                    raise IOError('Please specify a non-directory filepath for ModelCheckpoint. Filepath used is an existing directory: {}'.format(filepath))
                raise e

    def _get_file_path(self, epoch, logs):
        """Returns the file path for checkpoint."""
        try:
            file_path = self.filepath.format(epoch=epoch + 1, **logs)
        except KeyError as e:
            raise KeyError('Failed to format this callback filepath: "{}". Reason: {}'.format(self.filepath, e))
        self._write_filepath = distributed_file_utils.write_filepath(file_path, self.model.distribute_strategy)
        return self._write_filepath

    def _maybe_remove_file(self):
        distributed_file_utils.remove_temp_dir_with_filepath(self._write_filepath, self.model.distribute_strategy)

    def _checkpoint_exists(self, filepath):
        """Returns whether the checkpoint `filepath` refers to exists."""
        if filepath.endswith('.h5'):
            return file_io.file_exists_v2(filepath)
        tf_saved_model_exists = file_io.file_exists_v2(filepath)
        tf_weights_only_checkpoint_exists = file_io.file_exists_v2(filepath + '.index')
        return tf_saved_model_exists or tf_weights_only_checkpoint_exists

    def _get_most_recently_modified_file_matching_pattern(self, pattern):
        """Returns the most recently modified filepath matching pattern.

    Pattern may contain python formatting placeholder. If
    `tf.train.latest_checkpoint()` does not return None, use that; otherwise,
    check for most recently modified one that matches the pattern.

    In the rare case where there are more than one pattern-matching file having
    the same modified time that is most recent among all, return the filepath
    that is largest (by `>` operator, lexicographically using the numeric
    equivalents). This provides a tie-breaker when multiple files are most
    recent. Note that a larger `filepath` can sometimes indicate a later time of
    modification (for instance, when epoch/batch is used as formatting option),
    but not necessarily (when accuracy or loss is used). The tie-breaker is
    put in the logic as best effort to return the most recent, and to avoid
    undeterministic result.

    Modified time of a file is obtained with `os.path.getmtime()`.

    This utility function is best demonstrated via an example:

    ```python
    file_pattern = 'f.batch{batch:02d}epoch{epoch:02d}.h5'
    test_dir = self.get_temp_dir()
    path_pattern = os.path.join(test_dir, file_pattern)
    file_paths = [
        os.path.join(test_dir, file_name) for file_name in
        ['f.batch03epoch02.h5', 'f.batch02epoch02.h5', 'f.batch01epoch01.h5']
    ]
    for file_path in file_paths:
      # Write something to each of the files
    self.assertEqual(
        _get_most_recently_modified_file_matching_pattern(path_pattern),
        file_paths[-1])
    ```

    Args:
        pattern: The file pattern that may optionally contain python placeholder
            such as `{epoch:02d}`.

    Returns:
        The most recently modified file's full filepath matching `pattern`. If
        `pattern` does not contain any placeholder, this returns the filepath
        that
        exactly matches `pattern`. Returns `None` if no match is found.
    """
        dir_name = os.path.dirname(pattern)
        base_name = os.path.basename(pattern)
        base_name_regex = '^' + re.sub('{.*}', '.*', base_name) + '$'
        latest_tf_checkpoint = checkpoint_management.latest_checkpoint(dir_name)
        if latest_tf_checkpoint is not None and re.match(base_name_regex, os.path.basename(latest_tf_checkpoint)):
            return latest_tf_checkpoint
        latest_mod_time = 0
        file_path_with_latest_mod_time = None
        n_file_with_latest_mod_time = 0
        file_path_with_largest_file_name = None
        if file_io.file_exists_v2(dir_name):
            for file_name in os.listdir(dir_name):
                if re.match(base_name_regex, file_name):
                    file_path = os.path.join(dir_name, file_name)
                    mod_time = os.path.getmtime(file_path)
                    if file_path_with_largest_file_name is None or file_path > file_path_with_largest_file_name:
                        file_path_with_largest_file_name = file_path
                    if mod_time > latest_mod_time:
                        latest_mod_time = mod_time
                        file_path_with_latest_mod_time = file_path
                        n_file_with_latest_mod_time = 1
                    elif mod_time == latest_mod_time:
                        n_file_with_latest_mod_time += 1
        if n_file_with_latest_mod_time == 1:
            return file_path_with_latest_mod_time
        else:
            return file_path_with_largest_file_name

Callback to save the Keras model or model weights at some frequency.

`ModelCheckpoint` callback is used in conjunction with training using
`model.fit()` to save a model or weights (in a checkpoint file) at some
interval, so the model or weights can be loaded later to continue the training
from the state saved.

A few options this callback provides include:

- Whether to only keep the model that has achieved the "best performance" so
  far, or whether to save the model at the end of every epoch regardless of
  performance.
- Definition of 'best'; which quantity to monitor and whether it should be
  maximized or minimized.
- The frequency it should save at. Currently, the callback supports saving at
  the end of every epoch, or after a fixed number of training batches.
- Whether only weights are saved, or the whole model is saved.

Note: If you get `WARNING:tensorflow:Can save best model only with <name>
available, skipping` see the description of the `monitor` argument for
details on how to get this right.

Example:

```python
model.compile(loss=..., optimizer=...,
              metrics=['accuracy'])

EPOCHS = 10
checkpoint_filepath = '/tmp/checkpoint'
model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_filepath,
    save_weights_only=True,
    monitor='val_accuracy',
    mode='max',
    save_best_only=True)

# Model weights are saved at the end of every epoch, if it's the best seen
# so far.
model.fit(epochs=EPOCHS, callbacks=[model_checkpoint_callback])

# The model weights (that are considered the best) are loaded into the model.
model.load_weights(checkpoint_filepath)
```

Args:
    filepath: string or `PathLike`, path to save the model file. e.g.
      filepath = os.path.join(working_dir, 'ckpt', file_name). `filepath`
      can contain named formatting options, which will be filled the value of
      `epoch` and keys in `logs` (passed in `on_epoch_end`). For example: if
      `filepath` is `weights.{epoch:02d}-{val_loss:.2f}.hdf5`, then the model
      checkpoints will be saved with the epoch number and the validation loss
      in the filename. The directory of the filepath should not be reused by
      any other callbacks to avoid conflicts.
    monitor: The metric name to monitor. Typically the metrics are set by the
      `Model.compile` method. Note:

      * Prefix the name with `"val_`" to monitor validation metrics.
      * Use `"loss"` or "`val_loss`" to monitor the model's total loss.
      * If you specify metrics as strings, like `"accuracy"`, pass the same
        string (with or without the `"val_"` prefix).
      * If you pass `metrics.Metric` objects, `monitor` should be set to
        `metric.name`
      * If you're not sure about the metric names you can check the contents
        of the `history.history` dictionary returned by
        `history = model.fit()`
      * Multi-output models set additional prefixes on the metric names.

    verbose: verbosity mode, 0 or 1.
    save_best_only: if `save_best_only=True`, it only saves when the model
      is considered the "best" and the latest best model according to the
      quantity monitored will not be overwritten. If `filepath` doesn't
      contain formatting options like `{epoch}` then `filepath` will be
      overwritten by each new better model.
    mode: one of {'auto', 'min', 'max'}. If `save_best_only=True`, the
      decision to overwrite the current save file is made based on either
      the maximization or the minimization of the monitored quantity.
      For `val_acc`, this should be `max`, for `val_loss` this should be
      `min`, etc. In `auto` mode, the mode is set to `max` if the quantities
      monitored are 'acc' or start with 'fmeasure' and are set to `min` for
      the rest of the quantities.
    save_weights_only: if True, then only the model's weights will be saved
      (`model.save_weights(filepath)`), else the full model is saved
      (`model.save(filepath)`).
    save_freq: `'epoch'` or integer. When using `'epoch'`, the callback saves
      the model after each epoch. When using integer, the callback saves the
      model at end of this many batches. If the `Model` is compiled with
      `steps_per_execution=N`, then the saving criteria will be
      checked every Nth batch. Note that if the saving isn't aligned to
      epochs, the monitored metric may potentially be less reliable (it
      could reflect as little as 1 batch, since the metrics get reset every
      epoch). Defaults to `'epoch'`.
    options: Optional `tf.train.CheckpointOptions` object if
      `save_weights_only` is true or optional `tf.saved_model.SaveOptions`
      object if `save_weights_only` is false.
    **kwargs: Additional arguments for backwards compatibility. Possible key
      is `period`.


def _run_simple_loop_test(self, mode, inp, body, out):
    self._maybe_skip(mode)
    with ops.device(_get_device(mode)):
        random_seed.set_random_seed(0)
        expected_types = []
        for section in [inp, body, out]:
            section_expected_types = []
            for color in section:
                if color.isupper():
                    expected_type = self._lower_precision_dtype(mode).as_datatype_enum
                else:
                    expected_type = types_pb2.DT_FLOAT
                section_expected_types.append(expected_type)
            expected_types.append(section_expected_types)
        a = _build_simple_loop_graph(inp, body, out)
    output_val_ref, output_val, cost_graph = self._run(mode, a)
    node_map = _build_node_map(cost_graph.node)
    section_names = ['input', 'while/body', 'output']
    all_types_correct = True
    for section_name, expected_types in zip(section_names, expected_types):
        for i, expected_type in enumerate(expected_types):
            node_name = section_name + '_%i' % i
            output_port = 0
            optimized_type = node_map[node_name].output_info[output_port].dtype
            if optimized_type != expected_type:
                print('Expected node %s to have type %s but got type %s' % (node_name, expected_type, optimized_type))
                all_types_correct = False
    self.assertTrue(all_types_correct)
    if mode == 'mkl':
        self.assertAllClose(output_val_ref, output_val, atol=0.02, rtol=0.02)
    else:
        self.assertAllClose(output_val_ref, output_val, atol=0.002, rtol=0.001)

Runs a test of a simple loop.

The loop has different node colors in different sections of the graph. The
arguments must be strings where each character represents the color of a
node in that section of the graph: w = allow, g = infer, c = clear,
b = deny. CAPITALIZED characters indicate that the node is expected to be
changed to DT_HALF during graph optimization.

inp -> loop [ body ] -> out.

Args:
  mode: Either 'cuda' or 'mkl'.
  inp: A string of letters indicating the colors and expected dtypes of the
    input nodes.
  body: A string of letters indicating the colors and expected dtypes of the
    body nodes.
  out: A string of letters indicating the colors and expected dtypes of the
    output nodes.


def upload_part(self, request):
    try:
        boto_response = self.client.upload_part(Body=request.bytes, Bucket=request.bucket, Key=request.object, PartNumber=request.part_number, UploadId=request.upload_id)
        response = messages.UploadPartResponse(boto_response['ETag'], request.part_number)
        return response
    except Exception as e:
        raise messages.S3ClientError(str(e), get_http_error_code(e))

Uploads part of a file to S3 during a multipart upload

Args:
  request: (UploadPartRequest) input message
Returns:
  (UploadPartResponse) The response message.


def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
    q_embed = q * cos + rotate_half(q) * sin
    k_embed = k * cos + rotate_half(k) * sin
    return (q_embed, k_embed)

Applies Rotary Position Embedding to the query and key tensors.

Args:
    q (`torch.Tensor`): The query tensor.
    k (`torch.Tensor`): The key tensor.
    cos (`torch.Tensor`): The cosine part of the rotary embedding.
    sin (`torch.Tensor`): The sine part of the rotary embedding.
    position_ids (`torch.Tensor`):
        The position indices of the tokens corresponding to the query and key tensors. For example, this can be
        used to pass offsetted position ids when working with a KV-cache.
    unsqueeze_dim (`int`, *optional*, defaults to 1):
        The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
        sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
        that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
        k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
        cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
        the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
Returns:
    `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.


def _trace_variant_creation(self):
    variant = self._variant_tensor
    if not isinstance(variant, ops.EagerTensor):
        raise NotImplementedError('Constructing a tf.function that reproduces a given dataset is only supported for datasets created eagerly. Please file a feature request if this is important to you.')
    with context.eager_mode(), ops.device('CPU'):
        graph_def = graph_pb2.GraphDef().FromString(self._as_serialized_graph(external_state_policy=options_lib.ExternalStatePolicy.FAIL).numpy())
    output_node_names = []
    for node in graph_def.node:
        if node.op == '_Retval':
            output_node_names = node.input
    if len(output_node_names) != 1:
        raise AssertionError(f'Dataset graph is expected to only have one return value but found {len(output_node_names)} return values: {output_node_names}.')
    output_node_name = output_node_names[0]
    file_path_nodes = {}
    if ops.get_default_graph().building_function:
        asset_tracker = self._maybe_track_assets(graph_def)
        for key in asset_tracker:
            assets_list = [array_ops.expand_dims(asset.asset_path, axis=0) for asset in asset_tracker[key]]
            file_path_nodes[key] = array_ops.concat(assets_list, axis=0)
    variant_function = wrap_function.function_from_graph_def(graph_def, inputs=[], outputs=output_node_name + ':0', captures=file_path_nodes)
    for used_function in self._functions():
        used_function.function.add_to_graph(variant_function.graph)
    return variant_function

Traces a function which outputs a variant `tf.Tensor` for this dataset.

Note that creating this function involves evaluating an op, and is currently
only supported when executing eagerly.

Returns:
  A zero-argument `ConcreteFunction` which outputs a variant `tf.Tensor`.


def _export_to_saved_model_graph(self, object_map, tensor_map, options, **kwargs):
    _, _, _ = (object_map, tensor_map, options)
    del kwargs
    return []

Creates a copy of this object's tensors onto SavedModel graph.

Needs to be overridden if the class contains tensors that must be saved
into the graph. This method should update the `object_map` and `tensor_map`
dictionaries.

This method is called on all nodes in the Trackable Graph (generated by
`_trackable_children`). The nodes are traversed in the order defined by
`_deserialization_dependencies`

All usages of _map_resources should be migrated to this method.

Args:
  object_map: A dictionary that maps original Trackables to the copied
    Trackables. This only needs to be updated if the object is a
    tf.function, or if the copied tensors are necessary for checkpointing
    this object.
  tensor_map: Dictionary mapping original tensors to copied tensors.
  options: A `tf.saved_model.SaveOptions` object.
  **kwargs: Additional kwargs that may be added at a later time.

Returns:
  Flat list of original tensors that have been copied.


class BasicTokenizer:

    def __init__(self, do_lower_case=True, never_split=None, tokenize_chinese_chars=True, strip_accents=None):
        if never_split is None:
            never_split = []
        self.do_lower_case = do_lower_case
        self.never_split = set(never_split)
        self.tokenize_chinese_chars = tokenize_chinese_chars
        self.strip_accents = strip_accents

    def tokenize(self, text, never_split=None):
        """
        Basic Tokenization of a piece of text. Split on "white spaces" only, for sub-word tokenization, see
        WordPieceTokenizer.

        Args:
            **never_split**: (*optional*) list of str
                Kept for backward compatibility purposes. Now implemented directly at the base class level (see
                [`PreTrainedTokenizer.tokenize`]) List of token not to split.
        """
        never_split = self.never_split.union(set(never_split)) if never_split else self.never_split
        text = self._clean_text(text)
        if self.tokenize_chinese_chars:
            text = self._tokenize_chinese_chars(text)
        orig_tokens = whitespace_tokenize(text)
        split_tokens = []
        for token in orig_tokens:
            if token not in never_split:
                if self.do_lower_case:
                    token = token.lower()
                    if self.strip_accents is not False:
                        token = self._run_strip_accents(token)
                elif self.strip_accents:
                    token = self._run_strip_accents(token)
            split_tokens.extend(self._run_split_on_punc(token, never_split))
        output_tokens = whitespace_tokenize(' '.join(split_tokens))
        return output_tokens

    def _run_strip_accents(self, text):
        """Strips accents from a piece of text."""
        text = unicodedata.normalize('NFD', text)
        output = []
        for char in text:
            cat = unicodedata.category(char)
            if cat == 'Mn':
                continue
            output.append(char)
        return ''.join(output)

    def _run_split_on_punc(self, text, never_split=None):
        """Splits punctuation on a piece of text."""
        if never_split is not None and text in never_split:
            return [text]
        chars = list(text)
        i = 0
        start_new_word = True
        output = []
        while i < len(chars):
            char = chars[i]
            if _is_punctuation(char):
                output.append([char])
                start_new_word = True
            else:
                if start_new_word:
                    output.append([])
                start_new_word = False
                output[-1].append(char)
            i += 1
        return [''.join(x) for x in output]

    def _tokenize_chinese_chars(self, text):
        """Adds whitespace around any CJK character."""
        output = []
        for char in text:
            cp = ord(char)
            if self._is_chinese_char(cp):
                output.append(' ')
                output.append(char)
                output.append(' ')
            else:
                output.append(char)
        return ''.join(output)

    def _is_chinese_char(self, cp):
        """Checks whether CP is the codepoint of a CJK character."""
        if cp >= 19968 and cp <= 40959 or (cp >= 13312 and cp <= 19903) or (cp >= 131072 and cp <= 173791) or (cp >= 173824 and cp <= 177983) or (cp >= 177984 and cp <= 178207) or (cp >= 178208 and cp <= 183983) or (cp >= 63744 and cp <= 64255) or (cp >= 194560 and cp <= 195103):
            return True
        return False

    def _clean_text(self, text):
        """Performs invalid character removal and whitespace cleanup on text."""
        output = []
        for char in text:
            cp = ord(char)
            if cp == 0 or cp == 65533 or _is_control(char):
                continue
            if _is_whitespace(char):
                output.append(' ')
            else:
                output.append(char)
        return ''.join(output)

Constructs a BasicTokenizer that will run basic tokenization (punctuation splitting, lower casing, etc.).

Args:
    do_lower_case (`bool`, *optional*, defaults to `True`):
        Whether or not to lowercase the input when tokenizing.
    never_split (`Iterable`, *optional*):
        Collection of tokens which will never be split during tokenization. Only has an effect when
        `do_basic_tokenize=True`
    tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
        Whether or not to tokenize Chinese characters.

        This should likely be deactivated for Japanese (see this
        [issue](https://github.com/huggingface/transformers/issues/328)).
    strip_accents (`bool`, *optional*):
        Whether or not to strip all accents. If this option is not specified, then it will be determined by the
        value for `lowercase` (as in the original BERT).


class FlaxBaseModelOutputWithPastAndCrossAttentions(ModelOutput):
    last_hidden_state: Optional[jnp.ndarray] = None
    past_key_values: Optional[Tuple[Tuple[jnp.ndarray]]] = None
    hidden_states: Optional[Tuple[jnp.ndarray]] = None
    attentions: Optional[Tuple[jnp.ndarray]] = None
    cross_attentions: Optional[Tuple[jnp.ndarray]] = None

Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding).

Args:
    last_hidden_state (`jnp.ndarray` of shape `(batch_size, sequence_length, hidden_size)`):
        Sequence of hidden-states at the output of the last layer of the model.

        If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1,
        hidden_size)` is output.
    past_key_values (`tuple(tuple(jnp.ndarray))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
        Tuple of `tuple(jnp.ndarray)` of length `config.n_layers`, with each tuple having 2 tensors of shape
        `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and optionally if
        `config.is_encoder_decoder=True` 2 additional tensors of shape `(batch_size, num_heads,
        encoder_sequence_length, embed_size_per_head)`.

        Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if
        `config.is_encoder_decoder=True` in the cross-attention blocks) that can be used (see `past_key_values`
        input) to speed up sequential decoding.
    hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape
        `(batch_size, sequence_length, hidden_size)`.

        Hidden-states of the model at the output of each layer plus the initial embedding outputs.
    attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
        sequence_length)`.

        Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
        heads.
    cross_attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`):
        Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
        sequence_length)`.

        Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the
        weighted average in the cross-attention heads.


class UdopProcessor(ProcessorMixin):
    attributes = ['image_processor', 'tokenizer']
    image_processor_class = 'LayoutLMv3ImageProcessor'
    tokenizer_class = ('UdopTokenizer', 'UdopTokenizerFast')
    optional_call_args = ['text_pair']

    def __init__(self, image_processor, tokenizer):
        super().__init__(image_processor, tokenizer)

    def __call__(self, images: Optional[ImageInput]=None, text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]=None, *args, audio=None, videos=None, **kwargs: Unpack[UdopProcessorKwargs]) -> BatchFeature:
        """
        This method first forwards the `images` argument to [`~UdopImageProcessor.__call__`]. In case
        [`UdopImageProcessor`] was initialized with `apply_ocr` set to `True`, it passes the obtained words and
        bounding boxes along with the additional arguments to [`~UdopTokenizer.__call__`] and returns the output,
        together with the prepared `pixel_values`. In case [`UdopImageProcessor`] was initialized with `apply_ocr` set
        to `False`, it passes the words (`text`/``text_pair`) and `boxes` specified by the user along with the
        additional arguments to [`~UdopTokenizer.__call__`] and returns the output, together with the prepared
        `pixel_values`.

        Alternatively, one can pass `text_target` and `text_pair_target` to prepare the targets of UDOP.

        Please refer to the docstring of the above two methods for more information.
        """
        output_kwargs = self._merge_kwargs(UdopProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs, **self.prepare_and_validate_optional_call_args(*args))
        boxes = output_kwargs['text_kwargs'].pop('boxes', None)
        word_labels = output_kwargs['text_kwargs'].pop('word_labels', None)
        text_pair = output_kwargs['text_kwargs'].pop('text_pair', None)
        return_overflowing_tokens = output_kwargs['text_kwargs'].get('return_overflowing_tokens', False)
        return_offsets_mapping = output_kwargs['text_kwargs'].get('return_offsets_mapping', False)
        text_target = output_kwargs['text_kwargs'].get('text_target', None)
        if self.image_processor.apply_ocr and boxes is not None:
            raise ValueError('You cannot provide bounding boxes if you initialized the image processor with apply_ocr set to True.')
        if self.image_processor.apply_ocr and word_labels is not None:
            raise ValueError('You cannot provide word labels if you initialized the image processor with apply_ocr set to True.')
        if return_overflowing_tokens and (not return_offsets_mapping):
            raise ValueError('You cannot return overflowing tokens without returning the offsets mapping.')
        if text_target is not None:
            return self.tokenizer(**output_kwargs['text_kwargs'])
        else:
            features = self.image_processor(images=images, **output_kwargs['images_kwargs'])
            features_words = features.pop('words', None)
            features_boxes = features.pop('boxes', None)
            output_kwargs['text_kwargs'].pop('text_target', None)
            output_kwargs['text_kwargs'].pop('text_pair_target', None)
            output_kwargs['text_kwargs']['text_pair'] = text_pair
            output_kwargs['text_kwargs']['boxes'] = boxes if boxes is not None else features_boxes
            output_kwargs['text_kwargs']['word_labels'] = word_labels
            if text is not None and self.image_processor.apply_ocr and (text_pair is None):
                if isinstance(text, str):
                    text = [text]
                output_kwargs['text_kwargs']['text_pair'] = features_words
            encoded_inputs = self.tokenizer(text=text if text is not None else features_words, **output_kwargs['text_kwargs'])
            if return_overflowing_tokens is True:
                features['pixel_values'] = self.get_overflowing_images(features['pixel_values'], encoded_inputs['overflow_to_sample_mapping'])
            features.update(encoded_inputs)
            return features

    def get_overflowing_images(self, images, overflow_to_sample_mapping):
        images_with_overflow = []
        for sample_idx in overflow_to_sample_mapping:
            images_with_overflow.append(images[sample_idx])
        if len(images_with_overflow) != len(overflow_to_sample_mapping):
            raise ValueError(f'Expected length of images to be the same as the length of `overflow_to_sample_mapping`, but got {len(images_with_overflow)} and {len(overflow_to_sample_mapping)}')
        return images_with_overflow

    def batch_decode(self, *args, **kwargs):
        """
        This method forwards all its arguments to PreTrainedTokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please
        refer to the docstring of this method for more information.
        """
        return self.tokenizer.batch_decode(*args, **kwargs)

    def decode(self, *args, **kwargs):
        """
        This method forwards all its arguments to PreTrainedTokenizer's [`~PreTrainedTokenizer.decode`]. Please refer
        to the docstring of this method for more information.
        """
        return self.tokenizer.decode(*args, **kwargs)

    @property
    def model_input_names(self):
        return ['pixel_values', 'input_ids', 'bbox', 'attention_mask']

Constructs a UDOP processor which combines a LayoutLMv3 image processor and a UDOP tokenizer into a single processor.

[`UdopProcessor`] offers all the functionalities you need to prepare data for the model.

It first uses [`LayoutLMv3ImageProcessor`] to resize, rescale and normalize document images, and optionally applies OCR
to get words and normalized bounding boxes. These are then provided to [`UdopTokenizer`] or [`UdopTokenizerFast`],
which turns the words and bounding boxes into token-level `input_ids`, `attention_mask`, `token_type_ids`, `bbox`.
Optionally, one can provide integer `word_labels`, which are turned into token-level `labels` for token
classification tasks (such as FUNSD, CORD).

Additionally, it also supports passing `text_target` and `text_pair_target` to the tokenizer, which can be used to
prepare labels for language modeling tasks.

Args:
    image_processor (`LayoutLMv3ImageProcessor`):
        An instance of [`LayoutLMv3ImageProcessor`]. The image processor is a required input.
    tokenizer (`UdopTokenizer` or `UdopTokenizerFast`):
        An instance of [`UdopTokenizer`] or [`UdopTokenizerFast`]. The tokenizer is a required input.


def torch_distributed_zero_first(local_rank: int):
    if local_rank not in [-1, 0]:
        dist.barrier()
    yield
    if local_rank == 0:
        dist.barrier()

Decorator to make all processes in distributed training wait for each local_master to do something.

Args:
    local_rank (`int`): The rank of the local process.


def get_container_image_from_options(pipeline_options):
    worker_options = pipeline_options.view_as(WorkerOptions)
    if worker_options.sdk_container_image:
        return worker_options.sdk_container_image
    container_repo = names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY
    image_name = '{repository}/beam_python{major}.{minor}_sdk'.format(repository=container_repo, major=sys.version_info[0], minor=sys.version_info[1])
    image_tag = _get_required_container_version()
    return image_name + ':' + image_tag

For internal use only; no backwards-compatibility guarantees.

Args:
  pipeline_options (PipelineOptions): A container for pipeline options.

Returns:
  str: Container image for remote execution.


def _single_column_cell_selection_loss(token_logits, column_logits, labels, cell_index, col_index, cell_mask):
    labels_per_column, _ = reduce_sum(tf.cast(labels, tf.float32), col_index)
    column_label = tf.argmax(labels_per_column, axis=-1, output_type=tf.int32)
    no_cell_selected = tf.equal(tf.reduce_max(labels_per_column, axis=-1), 0)
    column_label = tf.where(no_cell_selected, tf.zeros_like(column_label), column_label)
    column_dist = tfp.distributions.Categorical(logits=column_logits)
    column_loss_per_example = -column_dist.log_prob(column_label)
    logits_per_cell, _ = reduce_mean(token_logits, cell_index)
    labels_per_cell, labels_index = reduce_max(tf.cast(labels, tf.int32), cell_index)
    column_id_for_cells = cell_index.project_inner(labels_index).indices
    column_mask = tf.cast(tf.equal(column_id_for_cells, tf.expand_dims(column_label, axis=1)), tf.float32)
    cell_dist = tfp.distributions.Bernoulli(logits=logits_per_cell)
    cell_log_prob = cell_dist.log_prob(labels_per_cell)
    cell_loss = -tf.reduce_sum(cell_log_prob * column_mask * cell_mask, axis=1)
    cell_loss /= tf.reduce_sum(column_mask * cell_mask, axis=1) + EPSILON_ZERO_DIVISION
    selection_loss_per_example = column_loss_per_example
    selection_loss_per_example += tf.where(no_cell_selected, tf.zeros_like(selection_loss_per_example), cell_loss)
    selected_column_id = tf.argmax(column_logits, axis=-1, output_type=tf.int32)
    selected_column_mask = tf.cast(tf.equal(column_id_for_cells, tf.expand_dims(selected_column_id, axis=-1)), tf.float32)
    selected_column_mask = tf.where(tf.equal(column_id_for_cells, 0), tf.zeros_like(selected_column_mask), selected_column_mask)
    logits_per_cell += CLOSE_ENOUGH_TO_LOG_ZERO * (1.0 - cell_mask * selected_column_mask)
    logits = gather(logits_per_cell, cell_index)
    return (selection_loss_per_example, logits)

Computes the loss for cell selection constrained to a single column. The loss is a hierarchical log-likelihood. The
model first predicts a column and then selects cells within that column (conditioned on the column). Cells outside
the selected column are never selected.

Args:
    token_logits (`tf.Tensor` of shape `(batch_size, sequence_length)`):
        Tensor containing the logits per token.
    column_logits (`tf.Tensor` of shape `(batch_size, max_num_cols)`):
        Tensor containing the logits per column.
    labels (`tf.Tensor` of shape `(batch_size, sequence_length)`):
        Labels per token.
    cell_index (`ProductIndexMap`):
        Index that groups tokens into cells.
    col_index (`IndexMap`):
        Index that groups tokens into columns.
    cell_mask (`tf.Tensor` of shape `(batch_size, max_num_rows * max_num_cols)`):
        Mask for cells that exist in the table (i.e. that are not padding).

Returns:
    selection_loss_per_example (`tf.Tensor` of shape `(batch_size,)`): Loss for each example. logits (`tf.Tensor`
    of shape `(batch_size, sequence_length)`): New logits which are only allowed to select cells in a single
    column. Logits outside of the most likely column according to *column_logits* will be set to a very low value
    (such that the probabilities are 0).


def write(self, file_prefix: str) -> str:

Serializes proto to disk.

Args:
  file_prefix: string prefix of the filepath.

Returns:
  The actual path the proto is written to.


def enrichment_transform(pcoll, enrichment_handler: str, handler_config: dict[str, Any], timeout: Optional[float]=30):
    options.YamlOptions.check_enabled(pcoll.pipeline, 'Enrichment')
    if not Enrichment:
        raise ValueError(f"gcp dependencies not installed. Cannot use {enrichment_handler} handler. Please install using 'pip install apache-beam[gcp]'.")
    if enrichment_handler == 'FeastFeatureStore' and (not FeastFeatureStoreEnrichmentHandler):
        raise ValueError("FeastFeatureStore handler requires 'feast' package to be installed. " + "Please install using 'pip install feast[gcp]' and try again.")
    handler_map = {'BigQuery': BigQueryEnrichmentHandler, 'BigTable': BigTableEnrichmentHandler, 'FeastFeatureStore': FeastFeatureStoreEnrichmentHandler, 'VertexAIFeatureStore': VertexAIFeatureStoreEnrichmentHandler}
    if enrichment_handler not in handler_map:
        raise ValueError(f'Unknown enrichment source: {enrichment_handler}')
    handler = handler_map[enrichment_handler](**handler_config)
    return pcoll | Enrichment(source_handler=handler, timeout=timeout)

The Enrichment transform allows one to dynamically enhance elements in a
pipeline by performing key-value lookups against external services like
APIs or databases.

Example using BigTable: ::

    - type: Enrichment
      config:
        enrichment_handler: 'BigTable'
        handler_config:
          project_id: 'apache-beam-testing'
          instance_id: 'beam-test'
          table_id: 'bigtable-enrichment-test'
          row_key: 'product_id'
        timeout: 30

For more information on Enrichment, see the [Beam docs](
https://beam.apache.org/documentation/transforms/python/elementwise/enrichment/).

Args:
    enrichment_handler (str): Specifies the source from where data needs
      to be extracted into the pipeline for enriching data. One of
      "BigQuery", "BigTable", "FeastFeatureStore" or "VertexAIFeatureStore".
    handler_config (str): Specifies the parameters for the respective
      enrichment_handler in a YAML/JSON format. To see the full set of
      handler_config parameters, see their corresponding doc pages:

        - [BigQueryEnrichmentHandler](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment_handlers.bigquery.html#apache_beam.transforms.enrichment_handlers.bigquery.BigQueryEnrichmentHandler)
        - [BigTableEnrichmentHandler](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment_handlers.bigtable.html#apache_beam.transforms.enrichment_handlers.bigtable.BigTableEnrichmentHandler)
        - [FeastFeatureStoreEnrichmentHandler](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment_handlers.feast_feature_store.html#apache_beam.transforms.enrichment_handlers.feast_feature_store.FeastFeatureStoreEnrichmentHandler)
        - [VertexAIFeatureStoreEnrichmentHandler](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store.html#apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store.VertexAIFeatureStoreEnrichmentHandler)
    timeout (float): Timeout for source requests in seconds. Defaults to 30
      seconds.


def compute_gradients(self, *args, **kwargs):
    return self._opt.compute_gradients(*args, **kwargs)

Compute gradients of "loss" for the variables in "var_list".

This simply wraps the compute_gradients() from the real optimizer. The
gradients will be aggregated in the apply_gradients() so that user can
modify the gradients like clipping with per replica global norm if needed.
The global norm with aggregated gradients can be bad as one replica's huge
gradients can hurt the gradients from other replicas.

Args:
  *args: Arguments for compute_gradients().
  **kwargs: Keyword arguments for compute_gradients().

Returns:
  A list of (gradient, variable) pairs.


def get_added_vocab(self) -> dict[str, int]:
    return self._added_tokens_encoder

Returns the added tokens in the vocabulary as a dictionary of token to index. Results might be different from
the fast call because for now we always add the tokens even if they are already in the vocabulary. This is
something we should change.

Returns:
    `Dict[str, int]`: The added tokens.


def get_instance_type_parameter(self, name: str, node: 'cfg.CFGNode | None'=None):
    del name
    if node is None:
        node = self.ctx.root_node
    return self.ctx.new_unsolvable(node)

Get a cfg.Variable of the instance's values for the type parameter.

Treating self as an abstract.Instance, gets the variable of its values for
the given type parameter. For the real implementation, see
SimpleValue.get_instance_type_parameter.

Args:
  name: The name of the type parameter.
  node: Optionally, the current CFG node.

Returns:
  A Variable which may be empty.


def preprocess(self, images: ImageInput, do_resize: Optional[bool]=None, size: Optional[Dict[str, int]]=None, resample: PILImageResampling=None, do_center_crop: Optional[bool]=None, crop_size: Optional[Dict[str, int]]=None, do_rescale: Optional[bool]=None, rescale_factor: Optional[float]=None, do_normalize: Optional[bool]=None, image_mean: Optional[Union[float, List[float]]]=None, image_std: Optional[Union[float, List[float]]]=None, return_tensors: Optional[Union[TensorType, str]]=None, data_format: Union[str, ChannelDimension]=ChannelDimension.FIRST, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> BatchFeature:
    do_resize = do_resize if do_resize is not None else self.do_resize
    size = size if size is not None else self.size
    resample = resample if resample is not None else self.resample
    do_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop
    crop_size = crop_size if crop_size is not None else self.crop_size
    do_rescale = do_rescale if do_rescale is not None else self.do_rescale
    rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
    do_normalize = do_normalize if do_normalize is not None else self.do_normalize
    image_mean = image_mean if image_mean is not None else self.image_mean
    image_std = image_std if image_std is not None else self.image_std
    images = make_list_of_images(images)
    if not valid_images(images):
        raise ValueError('Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, torch.Tensor, tf.Tensor or jax.ndarray.')
    validate_preprocess_arguments(do_rescale=do_rescale, rescale_factor=rescale_factor, do_normalize=do_normalize, image_mean=image_mean, image_std=image_std, do_center_crop=do_center_crop, crop_size=crop_size, do_resize=do_resize, size=size, resample=resample)
    images = [to_numpy_array(image) for image in images]
    if do_rescale and is_scaled_image(images[0]):
        logger.warning_once('It looks like you are trying to rescale already rescaled images. If the input images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again.')
    if input_data_format is None:
        input_data_format = infer_channel_dimension_format(images[0])
    if do_resize:
        images = [self.resize(image, size=size, resample=resample, input_data_format=input_data_format) for image in images]
    if do_center_crop:
        images = [self.center_crop(image, crop_size=crop_size, input_data_format=input_data_format) for image in images]
    if do_rescale:
        images = [self.rescale(image, rescale_factor=rescale_factor, input_data_format=input_data_format) for image in images]
    if do_normalize:
        images = [self.normalize(image, mean=image_mean, std=image_std, input_data_format=input_data_format) for image in images]
    images = [to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images]
    encoded_inputs = BatchFeature(data={'pixel_values': images}, tensor_type=return_tensors)
    return encoded_inputs

Prepares an image or batch of images for the model.

Args:
    images (`ImageInput`):
        The image or batch of images to be prepared. Expects a single or batch of images with pixel values
        ranging from 0 to 255. If passing in images with pixel values between 0 and 1, set `do_rescale=False`.
    do_resize (`bool`, *optional*, defaults to `self.do_resize`):
        Whether or not to resize the input. If `True`, will resize the input to the size specified by `size`.
    size (`Dict[str, int]`, *optional*, defaults to `self.size`):
        The size to resize the input to. Only has an effect if `do_resize` is set to `True`.
    resample (`PILImageResampling`, *optional*, defaults to `self.resample`):
        The resampling filter to use when resizing the input. Only has an effect if `do_resize` is set to
        `True`.
    do_center_crop (`bool`, *optional*, defaults to `self.do_center_crop`):
        Whether or not to center crop the input. If `True`, will center crop the input to the size specified by
        `crop_size`.
    crop_size (`Dict[str, int]`, *optional*, defaults to `self.crop_size`):
        The size to center crop the input to. Only has an effect if `do_center_crop` is set to `True`.
    do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
        Whether or not to rescale the input. If `True`, will rescale the input by dividing it by
        `rescale_factor`.
    rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
        The factor to rescale the input by. Only has an effect if `do_rescale` is set to `True`.
    do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
        Whether or not to normalize the input. If `True`, will normalize the input by subtracting `image_mean`
        and dividing by `image_std`.
    image_mean (`Union[float, List[float]]`, *optional*, defaults to `self.image_mean`):
        The mean to subtract from the input when normalizing. Only has an effect if `do_normalize` is set to
        `True`.
    image_std (`Union[float, List[float]]`, *optional*, defaults to `self.image_std`):
        The standard deviation to divide the input by when normalizing. Only has an effect if `do_normalize` is
        set to `True`.
    return_tensors (`str` or `TensorType`, *optional*):
        The type of tensors to return. Can be one of:
        - Unset: Return a list of `np.ndarray`.
        - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
        - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
        - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
        - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
    data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
        The channel dimension format for the output image. Can be one of:
        - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
        - `ChannelDimension.LAST`: image in (height, width, num_channels) format.
        - Unset: defaults to the channel dimension format of the input image.
    input_data_format (`ChannelDimension` or `str`, *optional*):
        The channel dimension format for the input image. If unset, the channel dimension format is inferred
        from the input image. Can be one of:
        - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
        - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
        - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.


def swap_curve_fit(float_leg_start_times: List[types.RealTensor], float_leg_end_times: List[types.RealTensor], float_leg_daycount_fractions: List[types.RealTensor], fixed_leg_start_times: List[types.RealTensor], fixed_leg_end_times: List[types.RealTensor], fixed_leg_daycount_fractions: List[types.RealTensor], fixed_leg_cashflows: List[types.RealTensor], present_values: List[types.RealTensor], initial_curve_rates: types.RealTensor, present_values_settlement_times: List[types.RealTensor]=None, float_leg_discount_rates: List[types.RealTensor]=None, float_leg_discount_times: List[types.RealTensor]=None, fixed_leg_discount_rates: List[types.RealTensor]=None, fixed_leg_discount_times: List[types.RealTensor]=None, optimizer: Callable[..., Any]=None, curve_interpolator: Callable[..., types.RealTensor]=None, instrument_weights: types.RealTensor=None, curve_tolerance: types.RealTensor=1e-08, maximum_iterations: types.IntTensor=50, dtype: tf.DType=None, name: str=None) -> scc.SwapCurveBuilderResult:
    with tf.name_scope(name or 'swap_curve'):
        if optimizer is None:
            optimizer = optimizers.conjugate_gradient_minimize
        present_values = _convert_to_tensors(dtype, present_values, 'present_values')
        dtype = present_values[0].dtype
        if present_values_settlement_times is None:
            pv_settlement_times = [tf.zeros([], dtype=dtype) for pv in present_values]
        else:
            pv_settlement_times = _convert_to_tensors(dtype, present_values_settlement_times, 'pv_settlement_times')
        float_leg_start_times = _convert_to_tensors(dtype, float_leg_start_times, 'float_leg_start_times')
        float_leg_end_times = _convert_to_tensors(dtype, float_leg_end_times, 'float_leg_end_times')
        float_leg_daycount_fractions = _convert_to_tensors(dtype, float_leg_daycount_fractions, 'float_leg_daycount_fractions')
        fixed_leg_start_times = _convert_to_tensors(dtype, fixed_leg_start_times, 'fixed_leg_start_times')
        fixed_leg_end_times = _convert_to_tensors(dtype, fixed_leg_end_times, 'fixed_leg_end_times')
        fixed_leg_daycount_fractions = _convert_to_tensors(dtype, fixed_leg_daycount_fractions, 'fixed_leg_daycount_fractions')
        fixed_leg_cashflows = _convert_to_tensors(dtype, fixed_leg_cashflows, 'fixed_leg_cashflows')
        present_values = tf.stack(present_values, axis=-1)
        if instrument_weights is None:
            instrument_weights = _initialize_instrument_weights(float_leg_end_times, fixed_leg_end_times, dtype=dtype)
        else:
            instrument_weights = _convert_to_tensors(dtype, instrument_weights, 'instrument_weights')
        if curve_interpolator is None:

            def default_interpolator(xi, x, y):
                return linear.interpolate(xi, x, y, dtype=dtype)
            curve_interpolator = default_interpolator
        self_discounting_float_leg = False
        self_discounting_fixed_leg = False
        if float_leg_discount_rates is None and fixed_leg_discount_rates is None:
            self_discounting_float_leg = True
            self_discounting_fixed_leg = True
            float_leg_discount_rates = [0.0]
            float_leg_discount_times = [0.0]
            fixed_leg_discount_rates = [0.0]
            fixed_leg_discount_times = [0.0]
        elif fixed_leg_discount_rates is None:
            fixed_leg_discount_rates = float_leg_discount_rates
            fixed_leg_discount_times = float_leg_discount_times
        elif float_leg_discount_rates is None:
            self_discounting_float_leg = True
            float_leg_discount_rates = [0.0]
            float_leg_discount_times = [0.0]
        float_leg_discount_rates = _convert_to_tensors(dtype, float_leg_discount_rates, 'float_disc_rates')
        float_leg_discount_rates = tf.stack(float_leg_discount_rates, axis=-1)
        float_leg_discount_times = _convert_to_tensors(dtype, float_leg_discount_times, 'float_disc_times')
        float_leg_discount_times = tf.stack(float_leg_discount_times, axis=-1)
        fixed_leg_discount_rates = _convert_to_tensors(dtype, fixed_leg_discount_rates, 'fixed_disc_rates')
        fixed_leg_discount_rates = tf.stack(fixed_leg_discount_rates, axis=-1)
        fixed_leg_discount_times = _convert_to_tensors(dtype, fixed_leg_discount_times, 'fixed_disc_times')
        fixed_leg_discount_times = tf.stack(fixed_leg_discount_times, axis=-1)
        if initial_curve_rates is not None:
            initial_rates = tf.convert_to_tensor(initial_curve_rates, dtype=dtype, name='initial_rates')
        else:
            raise ValueError('Initial state of the curve is not specified.')
        return _build_swap_curve(float_leg_start_times, float_leg_end_times, float_leg_daycount_fractions, fixed_leg_start_times, fixed_leg_end_times, fixed_leg_cashflows, fixed_leg_daycount_fractions, float_leg_discount_rates, float_leg_discount_times, fixed_leg_discount_rates, fixed_leg_discount_times, self_discounting_float_leg, self_discounting_fixed_leg, present_values, pv_settlement_times, optimizer, curve_interpolator, initial_rates, instrument_weights, curve_tolerance, maximum_iterations)

Constructs the zero swap curve using optimization.

A zero swap curve is a function of time which gives the interest rate that
can be used to project forward rates at arbitrary `t` for the valuation of
interest rate securities.

Suppose we have a set of `N` Interest Rate Swaps (IRS) `S_i` with increasing
expiries whose market prices are known.
Suppose also that the `i`th IRS issues cashflows at times `T_{ij}` where
`1 <= j <= n_i` and `n_i` is the number of cashflows (including expiry)
for the `i`th swap.
Denote by `T_i` the time of final payment for the `i`th swap
(hence `T_i = T_{i,n_i}`). This function estimates a set of rates `r(T_i)`
such that when these rates are interpolated to all other cashflow times,
the computed value of the swaps matches the market value of the swaps
(within some tolerance). Rates at intermediate times are interpolated using
the user specified interpolation method (the default interpolation method
is linear interpolation on rates).

#### Example:

The following example illustrates the usage by building an implied swap curve
from four vanilla (fixed to float) LIBOR swaps.

```python

dtype = np.float64

# Next we will set up LIBOR reset and payment times for four spot starting
# swaps with maturities 1Y, 2Y, 3Y, 4Y. The LIBOR rate spans 6M.

float_leg_start_times = [
          np.array([0., 0.5], dtype=dtype),
          np.array([0., 0.5, 1., 1.5], dtype=dtype),
          np.array([0., 0.5, 1.0, 1.5, 2.0, 2.5], dtype=dtype),
          np.array([0., 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5], dtype=dtype)
      ]

float_leg_end_times = [
          np.array([0.5, 1.0], dtype=dtype),
          np.array([0.5, 1., 1.5, 2.0], dtype=dtype),
          np.array([0.5, 1.0, 1.5, 2.0, 2.5, 3.0], dtype=dtype),
          np.array([0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0], dtype=dtype)
      ]

# Next we will set up start and end times for semi-annual fixed coupons.

fixed_leg_start_times = [
          np.array([0., 0.5], dtype=dtype),
          np.array([0., 0.5, 1., 1.5], dtype=dtype),
          np.array([0., 0.5, 1.0, 1.5, 2.0, 2.5], dtype=dtype),
          np.array([0., 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5], dtype=dtype)
      ]

fixed_leg_end_times = [
          np.array([0.5, 1.0], dtype=dtype),
          np.array([0.5, 1., 1.5, 2.0], dtype=dtype),
          np.array([0.5, 1.0, 1.5, 2.0, 2.5, 3.0], dtype=dtype),
          np.array([0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0], dtype=dtype)
      ]

# Next setup a trivial daycount for floating and fixed legs.

float_leg_daycount = [
          np.array([0.5, 0.5], dtype=dtype),
          np.array([0.5, 0.5, 0.5, 0.5], dtype=dtype),
          np.array([0.5, 0.5, 0.5, 0.5, 0.5, 0.5], dtype=dtype),
          np.array([0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5], dtype=dtype)
      ]

fixed_leg_daycount = [
          np.array([0.5, 0.5], dtype=dtype),
          np.array([0.5, 0.5, 0.5, 0.5], dtype=dtype),
          np.array([0.5, 0.5, 0.5, 0.5, 0.5, 0.5], dtype=dtype),
          np.array([0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5], dtype=dtype)
      ]

fixed_leg_cashflows = [
      # 1 year swap with 2.855% semi-annual fixed payments.
      np.array([-0.02855, -0.02855], dtype=dtype),
      # 2 year swap with 3.097% semi-annual fixed payments.
      np.array([-0.03097, -0.03097, -0.03097, -0.03097], dtype=dtype),
      # 3 year swap with 3.1% semi-annual fixed payments.
      np.array([-0.031, -0.031, -0.031, -0.031, -0.031, -0.031], dtype=dtype),
      # 4 year swap with 3.2% semi-annual fixed payments.
      np.array([-0.032, -0.032, -0.032, -0.032, -0.032, -0.032, -0.032,
      -0.032], dtype=dtype)
  ]

# The present values of the above IRS.
pvs = np.array([0., 0., 0., 0.], dtype=dtype)

# Initial state of the curve.
initial_curve_rates = np.array([0.01, 0.01, 0.01, 0.01], dtype=dtype)

results = swap_curve_fit(float_leg_start_times, float_leg_end_times,
                         float_leg_daycount, fixed_leg_start_times,
                         fixed_leg_end_times, fixed_leg_cashflows,
                         fixed_leg_daycount, pvs, dtype=dtype,
                         initial_curve_rates=initial_curve_rates)

#### References:
[1]: Leif B.G. Andersen and Vladimir V. Piterbarg. Interest Rate Modeling,
    Volume I: Foundations and Vanilla Models. Chapter 6. 2010.

Args:
  float_leg_start_times: List of `Tensor`s. Each `Tensor` must be either of
    shape `batch_shape  + [k_i]` or `[k_i]` and of the same real dtype. `k_i`
    may be of different sizes. Each `Tensor` represents the beginning of the
    accrual period for the forward rate which determines the floating payment.
    Each element in the list belong to a unique swap to be used to build the
    curve.
  float_leg_end_times: List of `Tensor`s of shapes and `dtype` compatible with
    `float_leg_start_times`. Each `Tensor` represents the end of the
    accrual period for the forward rate which determines the floating payment.
  float_leg_daycount_fractions: List of `Tensor`s of shapes and `dtype`
    compatible with `float_leg_start_times`. Each `Tensor` represents the
    daycount fraction of the forward rate which determines the floating
    payment.
  fixed_leg_start_times: List of `Tensor`s. Each `Tensor` must be either of
    shape `batch_shape  + [n_i]` or `[n_i]` and of the same real dtype.
    `n_i` may be of different sizes. All elements must have the same `dtype`
    as `float_leg_start_times`. Each `Tensor` represents the beginning of the
    accrual period fixed coupon.
  fixed_leg_end_times: List of `Tensor`s of shapes and `dtype` compatible with
    `fixed_leg_start_times`. All elements must have the same `dtype` as
    `fixed_leg_start_times`. Each `Tensor` represents the
    end of the accrual period for the fixed coupon.
  fixed_leg_daycount_fractions: List of `Tensor`s of shapes and `dtype`
  compatible with
    `fixed_leg_start_times` Each `Tensor` represents the daycount fraction
    applicable for the fixed payment.
  fixed_leg_cashflows: List of `Tensor`s of shapes and `dtype` compatible with
    `fixed_leg_start_times`. The input contains fixed cashflows at each
    coupon payment time including notional (if any). The sign should be
    negative (positive) to indicate net outgoing (incoming) cashflow.
  present_values: List containing `Tensor`s of the same dtype as
    elements of `fixed_leg_cashflows` and of shapes compatible with
    `batch_shape`. The length of the list must be the same as the length of
    `fixed_leg_cashflows`. The input contains the market price of the
    underlying instruments.
  initial_curve_rates: A `Tensor` of the `dtype` as `present_values` and of
    shape `[batch_shape, num_instruments]` where `num_instruments` is the
    length of `float_leg_start_times`. The starting guess for the discount
    rates used to initialize the iterative procedure.
  present_values_settlement_times: Optional list of `Tensor`s with the shapes
    and `dtype` compatible with `present_values` The settlement times for the
    present values is the time from now when the instrument is traded to the
    time that the purchase price is actually delivered. If not supplied, then
    it is assumed that the settlement times are zero for every bond.
    Default value: `None` which is equivalent to zero settlement times.
  float_leg_discount_rates: Optional list of `Tensor`s with the shapes
    and `dtype` compatible with `present_values`. This input contains the
    continuously compounded discount rates the will be used
    to discount the floating cashflows. This allows the swap curve to
    constructed using an independent discount curve (e.g. OIS curve). By
    default the cashflows are discounted using the curve that is being
    constructed.
  float_leg_discount_times: Optional list of `Tensor`s with the shapes
    and `dtype` compatible with `present_values`. This input contains the
    times corresponding to the rates specified via
    the `float_leg_discount_rates`.
  fixed_leg_discount_rates: Optional list of `Tensor`s with the shapes
    and `dtype` compatible with `present_values`. This input contains the
    continuously compounded discount rates the will be used to discount the
    fixed cashflows. This allows the swap curve to constructed using an
    independent discount curve (e.g. OIS curve). By default the cashflows are
    discounted using the curve that is being constructed.
  fixed_leg_discount_times: Optional list of `Tensor`s with the shapes
    and `dtype` compatible with `present_values`. This input contains the
    times corresponding to the rates specified via the
    `fixed_leg_discount_rates`.
  optimizer: Optional Python callable which implements the algorithm used
    to minimize the objective function during calibration. It should have
    the following interface: result =
      optimizer(value_and_gradients_function, initial_position, tolerance,
      max_iterations) `value_and_gradients_function` is a Python callable that
      accepts a point as a real `Tensor` and returns a tuple of `Tensor`s of
      real dtype containing the value of the function and its gradient at that
      point. 'initial_position' is a real `Tensor` containing the starting
      point of the optimization, 'tolerance' is a real scalar `Tensor` for
      stopping tolerance for the procedure and `max_iterations` specifies the
      maximum number of iterations.
    `optimizer` should return a namedtuple containing the items: `position`
      (a tensor containing the optimal value), `converged` (a boolean
      indicating whether the optimize converged according the specified
      criteria), `failed` (a boolean indicating if the optimization resulted
      in a failure), `num_iterations` (the number of iterations used), and
      `objective_value` ( the value of the objective function at the optimal
      value). The default value for `optimizer` is None and conjugate
      gradient algorithm is used.
    Default value: `None` - indicating conjugate gradient minimizer.
  curve_interpolator: Optional Python callable used to interpolate the zero
    swap rates at cashflow times. It should have the following interface:
    yi = curve_interpolator(xi, x, y)
    `x`, `y`, 'xi', 'yi' are all `Tensors` of real dtype. `x` and `y` are the
    sample points and values (respectively) of the function to be
    interpolated. `xi` are the points at which the interpolation is
    desired and `yi` are the corresponding interpolated values returned by the
    function. The default value for `curve_interpolator` is None in which
    case linear interpolation is used.
    Default value: `None`. If not supplied, the yields to maturity for the
      bonds is used as the initial value.
  instrument_weights: Optional 'Tensor' of the same dtype and shape as
    `initial_curve_rates`. This input contains the weight of each instrument
    in computing the objective function for the conjugate gradient
    optimization. By default the weights are set to be the inverse of
    maturities.
  curve_tolerance: Optional positive scalar `Tensor` of same dtype as
    elements of `bond_cashflows`. The absolute tolerance for terminating the
    iterations used to fit the rate curve. The iterations are stopped when the
    estimated discounts at the expiry times of the bond_cashflows change by a
    amount smaller than `discount_tolerance` in an iteration.
    Default value: 1e-8.
  maximum_iterations: Optional positive integer `Tensor`. The maximum number
    of iterations permitted when fitting the curve.
    Default value: 50.
  dtype: `tf.Dtype`. If supplied the dtype for the (elements of)
    `float_leg_start_times`, and `fixed_leg_start_times`.
    Default value: None which maps to the default dtype inferred by
    TensorFlow.
  name: Python str. The name to give to the ops created by this function.
    Default value: `None` which maps to 'swap_curve'.

Returns:
  curve_builder_result: An instance of `SwapCurveBuilderResult` containing the
    following attributes.
    times: Rank 1 real `Tensor`. Times for the computed discount rates. These
      are chosen to be the expiry times of the supplied cashflows.
    discount_rates: Rank 1 `Tensor` of the same dtype as `times`.
      The inferred discount rates.
    discount_factor: Rank 1 `Tensor` of the same dtype as `times`.
      The inferred discount factors.
    initial_discount_rates: Rank 1 `Tensor` of the same dtype as `times`. The
      initial guess for the discount rates.
    converged: Scalar boolean `Tensor`. Whether the procedure converged.
      The procedure is said to have converged when the maximum absolute
      difference in the discount factors from one iteration to the next falls
      below the `discount_tolerance`.
    failed: Scalar boolean `Tensor`. Whether the procedure failed. Procedure
      may fail either because a NaN value was encountered for the discount
      rates or the discount factors.
    iterations: Scalar int32 `Tensor`. Number of iterations performed.
    objective_value: Scalar real `Tensor`. The value of the ibjective function
      evaluated using the fitted swap curve.

Raises:
  ValueError: If the initial state of the curve is not
  supplied to the function.


def list_fastboot_devices():
    out = fastboot.FastbootProxy().devices()
    return parse_device_list(out)

List all android devices connected to the computer that are in in
fastboot mode. These are detected by fastboot.

This function doesn't raise any error if `fastboot` binary doesn't exist,
because `FastbootProxy` itself doesn't raise any error.

Returns:
  A list of android device serials. Empty if there's none.


def _copy_trackable_to_cpu(self, object_map):
    del object_map
    raise NotImplementedError('Need to implement _copy_trackable_to_cpu() if the Trackable requires AsyncCheckpoint support.')

Creates a copy of this object onto CPU, also copies values over.

Needs to be overridden if the `Trackable` requires AsyncCheckpoint support.
The method first checks whether a copy of `self` is already created in
`object_map`, and creates one if not already created. Then the method copies
the **values** of itself over to its copy mapped by `object_map`.

Args:
  object_map: A dictionary that maps original Trackables to the copied
    Trackables, which reside in the CPU.


def dtype(self) -> torch.dtype:
    if self._rot_mats is not None:
        return self._rot_mats.dtype
    elif self._quats is not None:
        return self._quats.dtype
    else:
        raise ValueError('Both rotations are None')

Returns the dtype of the underlying rotation.

Returns:
    The dtype of the underlying rotation


def call(self, input_ids: TFModelInputType | None=None, attention_mask: np.ndarray | tf.Tensor | None=None, head_mask: np.ndarray | tf.Tensor | None=None, inputs_embeds: np.ndarray | tf.Tensor | None=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, training: Optional[bool]=False) -> Union[Tuple, TFBaseModelOutput]:
    encoder_outputs = self.encoder(input_ids, attention_mask=attention_mask, encoder_hidden_states=None, encoder_attention_mask=None, inputs_embeds=inputs_embeds, head_mask=head_mask, past_key_values=None, use_cache=False, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, training=training)
    if not return_dict:
        return encoder_outputs
    return TFBaseModelOutput(last_hidden_state=encoder_outputs.last_hidden_state, hidden_states=encoder_outputs.hidden_states, attentions=encoder_outputs.attentions)

Returns:

Examples:

```python
>>> from transformers import AutoTokenizer, TFT5EncoderModel

>>> tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-small")
>>> model = TFT5EncoderModel.from_pretrained("google-t5/t5-small")

>>> input_ids = tokenizer(
...     "Studies have been shown that owning a dog is good for you", return_tensors="tf"
... ).input_ids  # Batch size 1
>>> outputs = model(input_ids)
```


def generate_matches(patterns, nodes):
    if not patterns:
        yield (0, {})
    else:
        p, rest = (patterns[0], patterns[1:])
        for c0, r0 in p.generate_matches(nodes):
            if not rest:
                yield (c0, r0)
            else:
                for c1, r1 in generate_matches(rest, nodes[c0:]):
                    r = {}
                    r.update(r0)
                    r.update(r1)
                    yield (c0 + c1, r)

Generator yielding matches for a sequence of patterns and nodes.

Args:
    patterns: a sequence of patterns
    nodes: a sequence of nodes

Yields:
    (count, results) tuples where:
    count: the entire sequence of patterns matches nodes[:count];
    results: dict containing named submatches.


def run(code: str, *, global_vars: Optional[Dict[str, Any]]=None, permission: Optional[permissions.CodePermission]=None, returns_stdout: bool=False, outputs_intermediate: bool=False, sandbox: Optional[bool]=None, timeout: Optional[float]=None) -> Union[Any, Dict[str, Any]]:
    return maybe_sandbox_call(evaluate, code=code, global_vars=global_vars, permission=permission, returns_stdout=returns_stdout, outputs_intermediate=outputs_intermediate, sandbox=sandbox, timeout=timeout)

Executes Python code.

Features:
  * Fine-grained execution policy for limiting what APIs could be executed.
    This eliminates the need for sandboxing.
  * It exposes both the final results and intermediate results (variables).

Args:
  code: Python code to run.
  global_vars: An optional dict of
  permission: Permission for the Python code to run.
  returns_stdout: If True, the stdout (a str) will be returned.
  outputs_intermediate: Applicable when returns_stdout is False. If True,
    intermediate output will be outputted as a dict, with the last line's
    value accessible by key '__result__' and the std output accessible by
    key '__stdout__'. Otherwise the value of the last line will be returned.
  sandbox: If True, run code in sandbox; If False, run code in current
    process. If None, run in sandbox first, if the output could not be
    serialized and pass to current process, run the code again in current
    process.
  timeout: Execution timeout in seconds. If None, wait the code the complete.

Returns:
  The value of the last line of the code block. Or a dict of variable
  names of all locals to their evaluated values as the output of the code to
  run. The value for the last line can be accessed by key '__result__'. Or the
  stdout as a str.

Raises:
  TimeoutError: If the execution time exceeds the timeout.
  Exception: Exception  that are raised from the code.


def __init__(self, columns: list[str], bucket_boundaries: Iterable[Union[int, float]], name: Optional[str]=None):
    super().__init__(columns)
    self.bucket_boundaries = [bucket_boundaries]
    self.name = name

Interpolates values within the provided buckets and then normalizes to
[0, 1].

Input values are bucketized based on the provided boundaries such that the
input is mapped to a positive index i for which `bucket_boundaries[i-1] <=
element < bucket_boundaries[i]`, if it exists. The values are then
normalized to the range [0,1] within the bucket, with NaN values being
mapped to 0.5.

For more information, see:
https://www.tensorflow.org/tfx/transform/api_docs/python/tft/apply_buckets_with_interpolation

Args:
  columns: A list of column names to apply the transformation on.
  bucket_boundaries: An iterable of ints or floats representing the bucket
    boundaries sorted in ascending order.
  name: (Optional) A string that specifies the name of the operation.


def sample_paths(self, times, num_samples=1, initial_state=None, random_type=None, seed=None, **kwargs):
    pass

Returns a sample of paths from the process.

Args:
  times: A `Tensor` of positive real values of a shape `[T, k]`, where
    `T` is either empty or a shape which is broadcastable to `batch_shape`
    (as defined when this ItoProcess was initialised) and `k` is the
    number of time points. The definition of `batch_shape` and how it is set
    is dictated by the individual child classes. Typically `batch_shape`
    is set based on the shapes of the input parameters that define the
    ItoProcess, see `GeometricBrownianMotion` for an example. The times at
    which the path points are to be evaluated.
  num_samples: Positive scalar `int`. The number of paths to draw.
  initial_state: `Tensor` of shape `[dim]`. The initial state of the
    process.
    Default value: None which maps to a zero initial state.
  random_type: Enum value of `RandomType`. The type of (quasi)-random number
    generator to use to generate the paths.
    Default value: None which maps to the standard pseudo-random numbers.
  seed: Seed for the random number generator. The seed is
    only relevant if `random_type` is one of
    `[STATELESS, PSEUDO, HALTON_RANDOMIZED, PSEUDO_ANTITHETIC,
      STATELESS_ANTITHETIC]`. For `PSEUDO`, `PSEUDO_ANTITHETIC` and
    `HALTON_RANDOMIZED` the seed should be an Python integer. For
    `STATELESS` and  `STATELESS_ANTITHETIC `must be supplied as an integer
    `Tensor` of shape `[2]`.
    Default value: `None` which means no seed is set.
  **kwargs: Any other keyword args needed by an implementation.

Returns:
 A real `Tensor` of shape [`batch_shape`, num_samples, k, n] where `k` is
 the size of `times`, `n` is the dimension of the process and `batch_shape`
 is the batch dimensions of the ItoProcess.


def _post_process(self, feed_item, item):
    if item['assetIdentifier']['name']:
        feed_item[FieldMap.CREATIVE_ASSET_NAME] = item['assetIdentifier']['name']

Maps ids and names of related entities so they can be updated in the Bulkdozer feed.

When Bulkdozer is done processing an item, it writes back the updated names
and ids of related objects, this method makes sure those are updated in the
creative asset feed.

Args:
  feed_item: Feed item representing the creative asset from the Bulkdozer
    feed.
  item: The DCM creative asset being updated or created.


class Phi4MultimodalAudioConfig(PretrainedConfig):
    model_type = 'phi4_multimodal_audio'

    def __init__(self, hidden_size: int=1024, intermediate_size: int=1536, num_blocks: int=24, num_attention_heads: int=16, activation: str='swish', chunk_size: int=-1, left_chunk: int=18, dropout_rate: float=0.0, ext_pw_out_channel: int=1024, depthwise_seperable_out_channel: int=1024, depthwise_multiplier: int=1, kernel_size: int=3, conv_activation: str='swish', input_size: int=80, conv_glu_type: str='swish', time_reduction: int=8, bias_max_distance: int=1000, bias_symmetric: bool=False, nemo_activation: str='relu', nemo_conv_channels: int=1024, downsample_rate: int=1, initializer_range: float=0.02, audio_token_id: int=200011, feature_layer: int=-2, **kwargs):
        super().__init__(**kwargs)
        self.hidden_size = hidden_size
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.activation = activation
        self.chunk_size = chunk_size
        self.left_chunk = left_chunk
        self.num_blocks = num_blocks
        self.dropout_rate = dropout_rate
        self.ext_pw_out_channel = ext_pw_out_channel
        self.depthwise_seperable_out_channel = depthwise_seperable_out_channel
        self.depthwise_multiplier = depthwise_multiplier
        self.kernel_size = kernel_size
        self.conv_activation = conv_activation
        self.input_size = input_size
        self.conv_glu_type = conv_glu_type
        self.time_reduction = time_reduction
        self.bias_max_distance = bias_max_distance
        self.bias_symmetric = bias_symmetric
        self.nemo_activation = nemo_activation
        self.nemo_conv_channels = nemo_conv_channels
        self.downsample_rate = downsample_rate
        self.audio_token_id = audio_token_id
        self.initializer_range = initializer_range
        self.feature_layer = feature_layer
        if time_reduction % 2 != 0:
            raise ValueError('`time_reduction` should be a multiple of 2!')
        length = input_size
        for _ in range(int(math.log(time_reduction, 2))):
            length = math.floor((length - 1) / 2 + 1)
        self.nemo_final_size = length

This is the configuration class to store the configuration of a [`Phi4MultimodalAudioModel`]. It is used to instantiate a
Phi4Multimodal audio encoder according to the specified arguments, defining the model architecture. Instantiating a
configuration with the defaults will yield a similar configuration to that of the audio encoder of
[microsoft/Phi-4-multimodal-instruct](https://huggingface.co/microsoft/Phi-4-multimodal-instruct) architecture.

Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.

Args:
    hidden_size (`int`, *optional*, defaults to 1024):
        Dimensionality of the encoder layers.
    intermediate_size (`int`, *optional*, defaults to 1536):
        Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
    num_blocks (`int`, *optional*, defaults to 24):
        Number of hidden layers in the Transformer encoder.
    num_attention_heads (`int`, *optional*, defaults to 16):
        Number of attention heads for each attention layer in the Transformer encoder.
    activation (`str`, *optional*, defaults to `"swish"`):
        The non-linear activation function in the MLPs.
    chunk_size (`int`, *optional*, defaults to -1):
        The chunk size to create the masks.
    left_chunk (`int`, *optional*, defaults to 18):
        The left chunk to create the masks.
    dropout_rate (`float`, *optional*, defaults to 0.0):
        The dropout ratio.
    ext_pw_out_channel (`int`, *optional*, defaults to 1024):
        Number of out channels in the point-wise conv modules.
    depthwise_seperable_out_channel (`int`, *optional*, defaults to 1024):
        Number of out channels in the depth-wise separable conv modules.
    depthwise_multiplier (`int`, *optional*, defaults to 1):
        Input size multiplier for the depth-wise separable conv modules.
    kernel_size (`int`, *optional*, defaults to 3):
        Kernel size for the depth-wise separable conv modules.
    conv_activation (`str`, *optional*, defaults to `"swish"`):
        The non-linear activation function in the conv modules.
    input_size (`int`, *optional*, defaults to 80):
        Input size for the audio model.
    conv_glu_type (`str`, *optional*, defaults to `"swish"`):
        The non-linear activation function in the point-wise conv modules.
    time_reduction (`int`, *optional*, defaults to 8):
        Time reduction (subsampling factor).
    bias_max_distance (`int`, *optional*, defaults to 1000):
        Max distance for the relative attention bias module.
    bias_symmetric (`bool`, *optional*, defaults to `False`):
        Whether the relative attention bias should be symmetric or not.
    nemo_activation (`str`, *optional*, defaults to `"relu"`):
        The non-linear activation function in the nemo conv modules.
    nemo_conv_channels (`int`, *optional*, defaults to 1024):
        Number of channels in the nemo conv modules.
    downsample_rate (`int`, *optional*, defaults to 1):
        Downsample rate for the audio feature extractor.
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    audio_token_id (`int`, *optional*, defaults to 200011):
        The audio token id.
    feature_layer (`int`, *optional*, defaults to -2):
        The index of the layer of the encoder from which to extract audio features.

Example:

```python
>>> from transformers import Phi4MultimodalAudioConfig

>>> # Initializing a Phi4MultimodalAudioConfig with microsoft/Phi-4-multimodal-instruct style configuration
>>> configuration = Phi4MultimodalAudioConfig()
```


def get_image_features(self, pixel_values, params: Optional[dict]=None, dropout_rng: jax.random.PRNGKey=None, train=False):
    pixel_values = jnp.transpose(pixel_values, (0, 2, 3, 1))
    rngs = {}
    if dropout_rng is not None:
        rngs['dropout'] = dropout_rng

    def _get_features(module, pixel_values, deterministic):
        vision_outputs = module.vision_model(pixel_values=pixel_values, deterministic=deterministic)
        pooled_output = vision_outputs[1]
        image_features = module.visual_projection(pooled_output)
        return image_features
    return self.module.apply({'params': params or self.params}, jnp.array(pixel_values, dtype=jnp.float32), not train, method=_get_features, rngs=rngs)

Args:
    pixel_values (`numpy.ndarray` of shape `(batch_size, num_channels, height, width)`):
        Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained
        using [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details.

Returns:
    image_features (`jnp.ndarray` of shape `(batch_size, output_dim`): The image embeddings obtained by
    applying the projection layer to the pooled output of [`FlaxCLIPVisionModel`]

Examples:

```python
>>> from PIL import Image
>>> import requests
>>> from transformers import AutoProcessor, FlaxCLIPModel

>>> model = FlaxCLIPModel.from_pretrained("openai/clip-vit-base-patch32")
>>> processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32")

>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
>>> image = Image.open(requests.get(url, stream=True).raw)

>>> inputs = processor(images=image, return_tensors="np")

>>> image_features = model.get_image_features(**inputs)
```


def approx_min_k(operand, k, reduction_dimension=-1, recall_target=0.95, reduction_input_size_override=-1, aggregate_to_topk=True, name=None):
    return gen_nn_ops.approx_top_k(operand, k=k, reduction_dimension=reduction_dimension, recall_target=recall_target, is_max_k=False, reduction_input_size_override=reduction_input_size_override, aggregate_to_topk=aggregate_to_topk, name=name)

Returns min `k` values and their indices of the input `operand` in an approximate manner.

See https://arxiv.org/abs/2206.14286 for the algorithm details. This op is
only optimized on TPU currently.

Args:
  operand : Array to search for min-k. Must be a floating number type.
  k : Specifies the number of min-k.
  reduction_dimension: Integer dimension along which to search. Default: -1.
  recall_target: Recall target for the approximation.
  reduction_input_size_override : When set to a positive value, it overrides
    the size determined by `operand[reduction_dim]` for evaluating the recall.
    This option is useful when the given `operand` is only a subset of the
    overall computation in SPMD or distributed pipelines, where the true input
    size cannot be deferred by the `operand` shape.
  aggregate_to_topk: When true, aggregates approximate results to top-k. When
    false, returns the approximate results. The number of the approximate
    results is implementation defined and is greater equals to the specified
    `k`.
  name: Optional name for the operation.

Returns:
  Tuple of two arrays. The arrays are the least `k` values and the
  corresponding indices along the `reduction_dimension` of the input
  `operand`.  The arrays' dimensions are the same as the input `operand`
  except for the `reduction_dimension`: when `aggregate_to_topk` is true,
  the reduction dimension is `k`; otherwise, it is greater equals to `k`
  where the size is implementation-defined.

We encourage users to wrap `approx_min_k` with jit. See the following example
for nearest neighbor search over the squared l2 distance:

>>> import tensorflow as tf
>>> @tf.function(jit_compile=True)
... def l2_ann(qy, db, half_db_norms, k=10, recall_target=0.95):
...   dists = half_db_norms - tf.einsum('ik,jk->ij', qy, db)
...   return tf.nn.approx_min_k(dists, k=k, recall_target=recall_target)
>>>
>>> qy = tf.random.uniform((256,128))
>>> db = tf.random.uniform((2048,128))
>>> half_db_norms = tf.norm(db, axis=1) / 2
>>> dists, neighbors = l2_ann(qy, db, half_db_norms)

In the example above, we compute `db_norms/2 - dot(qy, db^T)` instead of
`qy^2 - 2 dot(qy, db^T) + db^2` for performance reason. The former uses less
arithmetics and produces the same set of neighbors.


def _assert_float_dtype(dtype):
    dtype = dtypes.as_dtype(dtype)
    if not dtype.is_floating:
        raise ValueError(f'Argument `dtype` is expected to be floating point. Received: {dtype}.')
    return dtype

Validate and return floating point type based on `dtype`.

`dtype` must be a floating point type.

Args:
  dtype: The data type to validate.

Returns:
  Validated type.

Raises:
  ValueError: if `dtype` is not a floating point type.


def initialize_system_for_tpu_embedding(embedding_config: embedding_pb2.TPUEmbeddingConfiguration, job: Optional[Text]=None) -> ops.Operation:
    config_string = embedding_config.SerializeToString()
    with ops.device(_tpu_system_device_name(job)):
        return tpu_ops.configure_tpu_embedding(config=config_string)

Initializes a distributed TPU Embedding system for use with TensorFlow.

The following two are equivalent:
1. initialize_system() with embedding_config.
2. initialize_system() without embedding_config, then
   initialize_system_for_tpu_embedding().
initialize_system() should not be called with embedding_config if
initialize_system_for_tpu_embedding() is meant to be called later.

Args:
  embedding_config: a `TPUEmbeddingConfiguration` proto describing the desired
    configuration of the hardware embedding lookup tables.
  job: The job (the XXX in TensorFlow device specification /job:XXX) that
    contains the TPU devices that will be initialized. If job=None it is
    assumed there is only one job in the TensorFlow flock, and an error will
    be returned if this assumption does not hold.

Returns:
  A no-op.


def _get_quadratic_coeffs(normalized_prices, normalized_forwards, log_normalized_forwards, option_signs, polya_factor):
    q1 = normalized_forwards - 1
    q2 = normalized_forwards + 1
    r = 2 * normalized_prices - option_signs * q1
    f1 = tf.math.pow(normalized_forwards, -polya_factor)
    f2 = 1 / f1
    g1 = f1 * normalized_forwards
    g2 = 1 / g1
    a = tf.math.square(g1 - g2)
    h = tf.math.square(normalized_forwards)
    r2 = tf.math.square(r)
    b = 4 * (f1 + f2) - 2 * (g1 + g2) * (1 + h - r2) / normalized_forwards
    lnc1 = tf.math.log(4.0 * normalized_prices) + tf.math.log(normalized_prices - option_signs * q1)
    lnc2 = tf.math.log(q2 - r) + tf.math.log(q2 + r)
    lnc = lnc1 + lnc2 - 2.0 * log_normalized_forwards
    return (a, b, lnc)

Computes the coefficients of the quadratic in Stefanica-Radiocic method.

Computes the coefficients described in Table 3 in Ref [1].

Args:
  normalized_prices: `Tensor` of real dtype and any shape. The prices of the
    options to be inverted. Normalization means that the raw price is divided
    by the strike discounted to the present.
  normalized_forwards: `Tensor` or same dtype and shape as `normalized_prices`
    The forwards divided by the strike of the options.
  log_normalized_forwards: `Tensor` or same dtype and shape as
    `normalized_prices`. Log of the normalized forwards.
  option_signs: Real `Tensor` of same shape and dtype as `normalized_prices`.
    Should be +1 for a Call option and -1 for a put option.
  polya_factor: Scalar `Tensor` of same dtype as `normalized_prices`. This is
    the factor to use for approximating the normal CDF in a Polya-like
    expression. Polya approximation is (here `k` is the `polya_factor`) N(x) ~
    0.5 + sign(x) sqrt(1-e^(-k x^2)) with k = 2 / pi. However, it has been
    found that other values for `k` may be more accurate. The value that
    minimizes the absolute error over the range [-10, 10] is 0.62305051
    (approximately 5/8).

Returns:
  (A, B, ln(C)): A 3-tuple of coefficients in terms of which the approximate
    implied vol is calculated.


def _get_python_exes(python_version) -> Iterable[list[str]]:
    if python_version in _CUSTOM_PYTHON_EXES:
        yield [_path_to_custom_exe(_CUSTOM_PYTHON_EXES[python_version])]
        return
    for version in (utils.format_version(python_version), '3'):
        if sys.platform == 'win32':
            python_exe = ['py', f'-{version}']
        else:
            python_exe = [f'python{version}']
        yield python_exe

Find possible python executables to use.

Arguments:
  python_version: the version tuple (e.g. (3, 7))

Yields:
  The path to the executable


def resize(self, image: np.ndarray, size: Dict[str, int], anti_aliasing: bool=True, anti_aliasing_sigma=None, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None, **kwargs) -> np.ndarray:
    requires_backends(self, 'scipy')
    output_shape = (size['height'], size['width'])
    image = to_channel_dimension_format(image, ChannelDimension.LAST)
    image, output_shape = _preprocess_resize_output_shape(image, output_shape)
    input_shape = image.shape
    factors = np.divide(input_shape, output_shape)
    ndi_mode = 'mirror'
    cval = 0
    order = 1
    if anti_aliasing:
        if anti_aliasing_sigma is None:
            anti_aliasing_sigma = np.maximum(0, (factors - 1) / 2)
        else:
            anti_aliasing_sigma = np.atleast_1d(anti_aliasing_sigma) * np.ones_like(factors)
            if np.any(anti_aliasing_sigma < 0):
                raise ValueError('Anti-aliasing standard deviation must be greater than or equal to zero')
            elif np.any((anti_aliasing_sigma > 0) & (factors <= 1)):
                warnings.warn('Anti-aliasing standard deviation greater than zero but not down-sampling along all axes')
        filtered = ndi.gaussian_filter(image, anti_aliasing_sigma, cval=cval, mode=ndi_mode)
    else:
        filtered = image
    zoom_factors = [1 / f for f in factors]
    out = ndi.zoom(filtered, zoom_factors, order=order, mode=ndi_mode, cval=cval, grid_mode=True)
    image = _clip_warp_output(image, out)
    image = to_channel_dimension_format(image, input_data_format, ChannelDimension.LAST)
    image = to_channel_dimension_format(image, data_format, input_data_format) if data_format is not None else image
    return image

Resize an image as per the original implementation.

Args:
    image (`np.ndarray`):
        Image to resize.
    size (`Dict[str, int]`):
        Dictionary containing the height and width to resize the image to.
    anti_aliasing (`bool`, *optional*, defaults to `True`):
        Whether to apply anti-aliasing when downsampling the image.
    anti_aliasing_sigma (`float`, *optional*, defaults to `None`):
        Standard deviation for Gaussian kernel when downsampling the image. If `None`, it will be calculated
        automatically.
    data_format (`str` or `ChannelDimension`, *optional*):
        The channel dimension format of the image. If not provided, it will be the same as the input image.
    input_data_format (`ChannelDimension` or `str`, *optional*):
        The channel dimension format of the input image. If not provided, it will be inferred from the input
        image.


def _get_init_rng(self):
    return self.seed_generator.next()

Returns a JAX `PRNGKey` or structure of `PRNGKey`s to pass to `init_fn`.

By default, this returns a single `PRNGKey` retrieved by calling
`self.seed_generator.next()`. Override this to return a different
structure.

Returns:
    a JAX `PRNGKey` or structure of `PRNGKey`s that will be passed as
    the `rng` argument of `init_fn`.