def load_graph_from_args(pipeline_name: str, framework: str, model: str, tokenizer: Optional[str]=None, **models_kwargs) -> Pipeline:
    if tokenizer is None:
        tokenizer = model
    if framework == 'pt' and (not is_torch_available()):
        raise Exception('Cannot convert because PyTorch is not installed. Please install torch first.')
    if framework == 'tf' and (not is_tf_available()):
        raise Exception('Cannot convert because TF is not installed. Please install tensorflow first.')
    print(f'Loading pipeline (model: {model}, tokenizer: {tokenizer})')
    return pipeline(pipeline_name, model=model, tokenizer=tokenizer, framework=framework, model_kwargs=models_kwargs)

Convert the set of arguments provided through the CLI to an actual pipeline reference (tokenizer + model

Args:
    pipeline_name: The kind of pipeline to use (ner, question-answering, etc.)
    framework: The actual model to convert the pipeline from ("pt" or "tf")
    model: The model name which will be loaded by the pipeline
    tokenizer: The tokenizer name which will be loaded by the pipeline, default to the model's value

Returns: Pipeline object


def find_custom_args_with_details(file_content: str, custom_args_var_name: str) -> list[dict]:
    escaped_variable_name = re.escape(custom_args_var_name)
    regex_pattern = f'^\\s*({escaped_variable_name})\\s*=\\s*(r?\\"\\"\\")(.*?)(\\"\\"\\")'
    flags = re.MULTILINE | re.DOTALL
    match = re.search(regex_pattern, file_content, flags)
    if match:
        content = match.group(3).strip()
        return content
    return None

Find the given custom args variable in the file content and return its content.

Args:
    file_content: The string content of the Python file.
    custom_args_var_name: The name of the custom args variable.


def restore_or_initialize(self):
    if self._latest_checkpoint is not None:
        self._checkpoint.restore(self._latest_checkpoint)
        if self._checkpoint_interval is not None:
            self._last_checkpoint_step = _evaluate(self._step_counter)
        return self._latest_checkpoint
    if self._init_fn is not None:
        self._init_fn()
        logging.info('Customized initialization is done through the passed `init_fn`.')
    return None

Restore items in `checkpoint` from the latest checkpoint file.

This method will first try to restore from the most recent checkpoint in
`directory`. If no checkpoints exist in `directory`, and `init_fn` is
specified, this method will call `init_fn` to do customized
initialization. This can be used to support initialization from pretrained
models.

Note that unlike `tf.train.Checkpoint.restore()`, this method doesn't return
a load status object that users can run assertions on
(e.g. assert_consumed()). Thus to run assertions, users should directly use
`tf.train.Checkpoint.restore()` method.

Returns:
  The restored checkpoint path if the latest checkpoint is found and
  restored. Otherwise None.


def deserialize(config, custom_objects=None):
    return serialization_lib.deserialize_keras_object(config, module_objects=globals(), custom_objects=custom_objects, printable_module_name='decay')

Instantiates a `LearningRateSchedule` object from a serialized form.

Args:
    config: The serialized form of the `LearningRateSchedule`. Dictionary of
        the form {'class_name': str, 'config': dict}.
    custom_objects: A dictionary mapping class names (or function names) of
        custom (non-Keras) objects to class/functions.

Returns:
    A `LearningRateSchedule` object.

Example:

```python
# Configuration for PolynomialDecay
config = {
    'class_name': 'PolynomialDecay',
    'config': {'cycle': False,
        'decay_steps': 10000,
        'end_learning_rate': 0.01,
        'initial_learning_rate': 0.1,
        'name': None,
        'power': 0.5
    }
}
lr_schedule = keras.optimizers.schedules.deserialize(config)
```


def __init__(self, strategy):
    if not getattr(self, '_has_initialized', False):
        if not hasattr(strategy, '_is_parameter_server_strategy_v2'):
            raise ValueError('Only `tf.distribute.experimental.ParameterServerStrategy` is supported to work with `tf.distribute.experimental.coordinator.ClusterCoordinator` currently.')
        self._strategy = strategy
        self.strategy.extended._used_with_coordinator = True
        self._cluster = Cluster(strategy)
        self._has_initialized = True

Initialization of a `ClusterCoordinator` instance.

Args:
  strategy: a supported `tf.distribute.Strategy` object. Currently, only
    `tf.distribute.experimental.ParameterServerStrategy` is supported.

Raises:
  ValueError: if the strategy being used is not supported.


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, cache_position: Optional[torch.LongTensor]=None, logits_to_keep: Union[int, torch.Tensor]=0, **kwargs: Unpack[KwargsForCausalLM]) -> CausalLMOutputWithPast:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    outputs: BaseModelOutputWithPast = self.model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, cache_position=cache_position, **kwargs)
    hidden_states = outputs.last_hidden_state
    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
    logits = self.lm_head(hidden_states[:, slice_indices, :])
    logits = logits * self.logit_scale
    loss = None
    if labels is not None:
        loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
    return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions)

labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
    Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
    config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
    (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

Example:

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

>> model = CohereForCausalLM.from_pretrained("CohereForAI/c4ai-command-r-v01")
>> tokenizer = AutoTokenizer.from_pretrained("CohereForAI/c4ai-command-r-v01")

>> prompt = "Hey, are you conscious? Can you talk to me?"
>> inputs = tokenizer(prompt, return_tensors="pt")

>> # Generate
>> generate_ids = model.generate(inputs.input_ids, max_length=30)
>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
```


def restore_nodes(save_path, nodes_to_restore):
    if save_path is None:
        raise ValueError('save_path cannot be empty.')
    if not isinstance(nodes_to_restore, dict):
        raise ValueError('Expecting a dictionary of node_id to Trackable for nodes_to_restore.')
    ckpt_view = checkpoint_view.CheckpointView(save_path)
    ckpt_view_descendants = ckpt_view.descendants()
    for node_id, trackable in nodes_to_restore.items():
        if node_id not in ckpt_view_descendants or ckpt_view._object_graph_proto.nodes[node_id] is None:
            raise ValueError(f'The expected node_id: {node_id} to Trackable {trackable} to restore does not exist in the checkpoint.')
        if trackable is None or not isinstance(trackable, base.Trackable):
            raise ValueError(f'Expecting a valid Trackable to node_id: {node_id} but got trackable: {trackable}.')
    serialized_tensors = object_identity.ObjectIdentityDictionary()
    for node_id, current_trackable in nodes_to_restore.items():
        ckpt_contains_serialized_tensors = ckpt_view._object_graph_proto.nodes[node_id].attributes
        node = ckpt_view._object_graph_proto.nodes[node_id]
        trackable_has_serialize_to_tensor = saveable_object_util.trackable_has_serialize_to_tensor(current_trackable)
        if not trackable_has_serialize_to_tensor:
            if not node.attributes:
                if saveable_object_util.saveable_objects_from_trackable(current_trackable):
                    raise ValueError(f'Trackable {current_trackable} expects checkpointed values but checkpoint does not contain serialized tensors for node_id: {node_id}.')
                else:
                    continue
            object_names = object_identity.ObjectIdentityDictionary()
            object_names[current_trackable] = trackable_utils.extract_object_name(node.attributes[0].checkpoint_key)
            checkpoint_factory_map, _ = save_util_v1.get_checkpoint_factories_and_keys(object_names, None)
            saveable_objects = save_util_v1.generate_saveable_objects(checkpoint_factory_map)[0]
            if len(node.attributes) != len(saveable_objects):
                raise ValueError(f'Size for saveable_objects for Trackable: {len(saveable_objects)} did not match the size for serialized_tensors for checkpoint: {len(node.attributes)}.')
            current_trackable = saveable_object_util.SaveableCompatibilityConverter(current_trackable, saveable_objects)
        serialized_tensors[current_trackable] = current_trackable._serialize_to_tensors()
        trackable_expects_ckpted_value = bool(serialized_tensors[current_trackable])
        if trackable_expects_ckpted_value and (not ckpt_contains_serialized_tensors):
            raise ValueError(f'Trackable {current_trackable} expects checkpointed values but checkpoint does not contain serialized tensors for node_id: {node_id}.')
        if not trackable_expects_ckpted_value and ckpt_contains_serialized_tensors:
            raise ValueError(f'Trackable {current_trackable} does not expect checkpointed values but checkpoint contains serialized tensors: {ckpt_contains_serialized_tensors} for node_id: {node_id}.')
        if len(node.attributes) != len(serialized_tensors[current_trackable]):
            raise ValueError(f'Size for serialized_tensors for Trackable: {len(serialized_tensors[current_trackable])} did not match size for serialized_tensors for checkpoint: {len(node.attributes)}.')
        if not trackable_has_serialize_to_tensor:
            functional_saver.MultiDeviceSaver(serialized_tensors).restore(save_path)
        else:
            serialized_tensors_renamed = object_identity.ObjectIdentityDictionary()
            serialized_tensors_renamed[current_trackable] = {}
            for attribute in node.attributes:
                name = attribute.name
                checkpoint_key = attribute.checkpoint_key
                serialized_tensors_renamed[current_trackable][checkpoint_key] = serialized_tensors[current_trackable][name]
            functional_saver.MultiDeviceSaver(serialized_tensors_renamed).restore(save_path)

Restores nodes from a dict.

Requires that the `Trackable` Python object has been bound to an object
ID in the checkpoint.

Args:
  save_path: a string represents path to the checkpoint.
  nodes_to_restore: a dict maps `node_id` to `trackable` to be restored.


def __init__(self, learning_rate: Union[float, Callable[[], float]]=0.001, beta_1: float=0.9, beta_2: float=0.999, epsilon: float=1e-07, lazy_adam: bool=True, sum_inside_sqrt: bool=True, use_gradient_accumulation: bool=True, clip_weight_min: Optional[float]=None, clip_weight_max: Optional[float]=None, weight_decay_factor: Optional[float]=None, multiply_weight_decay_factor_by_learning_rate: Optional[bool]=None, slot_variable_creation_fn: Optional[SlotVarCreationFnType]=None, clipvalue: Optional[ClipValueType]=None, low_dimensional_packing_status: bool=False):
    super(Adam, self).__init__(learning_rate, use_gradient_accumulation, clip_weight_min, clip_weight_max, weight_decay_factor, multiply_weight_decay_factor_by_learning_rate, clipvalue, slot_variable_creation_fn, low_dimensional_packing_status)
    if beta_1 < 0.0 or beta_1 >= 1.0:
        raise ValueError(f'Argument `beta_1` must be >= 0 and < 1. Received: {beta_1}.')
    if beta_2 < 0.0 or beta_2 >= 1.0:
        raise ValueError(f'Argument `beta_2` must be >= 0 and < 1. Received: {beta_1}.')
    if epsilon <= 0.0:
        raise ValueError('epsilon must be positive; got {}.'.format(epsilon))
    if not use_gradient_accumulation and (not lazy_adam):
        raise ValueError('When disabling lazy Adam (`lazy_adam=False`), gradient accumulation must be used. Set `use_gradient_accumulation` to False.')
    self.beta_1 = beta_1
    self.beta_2 = beta_2
    self.epsilon = epsilon
    self.lazy_adam = lazy_adam
    self.sum_inside_sqrt = sum_inside_sqrt

Optimization parameters for Adam.

See 'tensorflow/core/protobuf/tpu/optimization_parameters.proto' for a
complete description of these parameters and their impacts on the optimizer
algorithm.

Args:
  learning_rate: The learning rate. It should be a floating point value or a
    callable taking no arguments for a dynamic learning rate.
  beta_1: A float value. The exponential decay rate for the 1st moment
    estimates.
  beta_2: A float value. The exponential decay rate for the 2nd moment
    estimates.
  epsilon: A small constant for numerical stability.
  lazy_adam: Use lazy Adam instead of Adam. Lazy Adam trains faster.
  sum_inside_sqrt: When this is true, the Adam update formula is changed
    from `m / (sqrt(v) + epsilon)` to `m / sqrt(v + epsilon**2)`. This
    option improves the performance of TPU training and is not expected to
    harm model quality.
  use_gradient_accumulation: Setting this to `False` makes embedding
    gradients calculation less accurate but faster.
  clip_weight_min: the minimum value to clip by; None means -infinity.
  clip_weight_max: the maximum value to clip by; None means +infinity.
  weight_decay_factor: amount of weight decay to apply; None means that the
    weights are not decayed.
  multiply_weight_decay_factor_by_learning_rate: if true,
    `weight_decay_factor` is multiplied by the current learning rate.
  slot_variable_creation_fn: If you wish do directly control the creation of
    the slot variables, set this to a callable taking three parameters: a
    table variable, a list of slot names to create for it, and a list of
    initializers. This function should return a dict with the slot names as
    keys and the created variables as values with types matching the table
    variable. When set to None (the default), uses the built-in variable
    creation.
  clipvalue: Controls clipping of the gradient. Set to either a single
    positive scalar value to get clipping or a tiple of scalar values (min,
    max) to set a separate maximum or minimum. If one of the two entries is
    None, then there will be no clipping that direction.
  low_dimensional_packing_status: Status of the low-dimensional embedding
    packing optimization controls whether to optimize the packing of
    1-dimensional, 2-dimensional, and 4-dimensional embedding tables in
    memory.


def dummy_inputs(self):
    if self.config.use_lang_emb and self.config.n_langs > 1:
        return {'input_ids': tf.constant(MULTIPLE_CHOICE_DUMMY_INPUTS, dtype=tf.int32), 'langs': tf.constant(MULTIPLE_CHOICE_DUMMY_INPUTS, dtype=tf.int32)}
    else:
        return {'input_ids': tf.constant(MULTIPLE_CHOICE_DUMMY_INPUTS, dtype=tf.int32)}

Dummy inputs to build the network.

Returns:
    tf.Tensor with dummy inputs


def __sub__(self, period_tensor):
    return self + periods.PeriodTensor(-period_tensor.quantity(), period_tensor.period_type())

Subtracts a tensor of periods.

When subtracting months or years, the resulting day of the month is
decreased to the largest valid value if necessary. E.g. 31.03.2020 - 1 month
= 29.02.2020, 29.02.2020 - 1 year = 28.02.2019.

Args:
  period_tensor: a PeriodTensor object broadcastable to the shape of "self".

Returns:
  The new instance of DateTensor.


def resize(self, image, size, resample=None, default_to_square=True, max_size=None):
    resample = resample if resample is not None else PILImageResampling.BILINEAR
    self._ensure_format_supported(image)
    if not isinstance(image, PIL.Image.Image):
        image = self.to_pil_image(image)
    if isinstance(size, list):
        size = tuple(size)
    if isinstance(size, int) or len(size) == 1:
        if default_to_square:
            size = (size, size) if isinstance(size, int) else (size[0], size[0])
        else:
            width, height = image.size
            short, long = (width, height) if width <= height else (height, width)
            requested_new_short = size if isinstance(size, int) else size[0]
            if short == requested_new_short:
                return image
            new_short, new_long = (requested_new_short, int(requested_new_short * long / short))
            if max_size is not None:
                if max_size <= requested_new_short:
                    raise ValueError(f'max_size = {max_size} must be strictly greater than the requested size for the smaller edge size = {size}')
                if new_long > max_size:
                    new_short, new_long = (int(max_size * new_short / new_long), max_size)
            size = (new_short, new_long) if width <= height else (new_long, new_short)
    return image.resize(size, resample=resample)

Resizes `image`. Enforces conversion of input to PIL.Image.

Args:
    image (`PIL.Image.Image` or `np.ndarray` or `torch.Tensor`):
        The image to resize.
    size (`int` or `Tuple[int, int]`):
        The size to use for resizing the image. If `size` is a sequence like (h, w), output size will be
        matched to this.

        If `size` is an int and `default_to_square` is `True`, then image will be resized to (size, size). If
        `size` is an int and `default_to_square` is `False`, then smaller edge of the image will be matched to
        this number. i.e, if height > width, then image will be rescaled to (size * height / width, size).
    resample (`int`, *optional*, defaults to `PILImageResampling.BILINEAR`):
        The filter to user for resampling.
    default_to_square (`bool`, *optional*, defaults to `True`):
        How to convert `size` when it is a single int. If set to `True`, the `size` will be converted to a
        square (`size`,`size`). If set to `False`, will replicate
        [`torchvision.transforms.Resize`](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.Resize)
        with support for resizing only the smallest edge and providing an optional `max_size`.
    max_size (`int`, *optional*, defaults to `None`):
        The maximum allowed for the longer edge of the resized image: if the longer edge of the image is
        greater than `max_size` after being resized according to `size`, then the image is resized again so
        that the longer edge is equal to `max_size`. As a result, `size` might be overruled, i.e the smaller
        edge may be shorter than `size`. Only used if `default_to_square` is `False`.

Returns:
    image: A resized `PIL.Image.Image`.


class MinMaxNorm(Constraint):

    def __init__(self, min_value=0.0, max_value=1.0, rate=1.0, axis=0):
        self.min_value = min_value
        self.max_value = max_value
        self.rate = rate
        self.axis = axis

    def __call__(self, w):
        w = backend.convert_to_tensor(w)
        norms = ops.sqrt(ops.sum(ops.square(w), axis=self.axis, keepdims=True))
        desired = self.rate * ops.clip(norms, self.min_value, self.max_value) + (1 - self.rate) * norms
        return w * (desired / (backend.epsilon() + norms))

    def get_config(self):
        return {'min_value': self.min_value, 'max_value': self.max_value, 'rate': self.rate, 'axis': self.axis}

MinMaxNorm weight constraint.

Constrains the weights incident to each hidden unit
to have the norm between a lower bound and an upper bound.

Args:
    min_value: the minimum norm for the incoming weights.
    max_value: the maximum norm for the incoming weights.
    rate: rate for enforcing the constraint: weights will be
        rescaled to yield
        `(1 - rate) * norm + rate * norm.clip(min_value, max_value)`.
        Effectively, this means that rate=1.0 stands for strict
        enforcement of the constraint, while rate<1.0 means that
        weights will be rescaled at each step to slowly move
        towards a value inside the desired interval.
    axis: integer, axis along which to calculate weight norms.
        For instance, in a `Dense` layer the weight matrix
        has shape `(input_dim, output_dim)`,
        set `axis` to `0` to constrain each weight vector
        of length `(input_dim,)`.
        In a `Conv2D` layer with `data_format="channels_last"`,
        the weight tensor has shape
        `(rows, cols, input_depth, output_depth)`,
        set `axis` to `[0, 1, 2]`
        to constrain the weights of each filter tensor of size
        `(rows, cols, input_depth)`.


def analyze(self) -> Sequence[_HasReturnT]:

Calls every signature of this function with appropriate fake arguments.

Returns:
  A sequence of objects with information about the result of calling the
  function with each of its signatures, with get_return_value() methods
  that retrieve the return values.


def post_process_semantic_segmentation(self, outputs, target_sizes: Optional[List[Tuple[int, int]]]=None) -> 'torch.Tensor':
    class_queries_logits = outputs.class_queries_logits
    masks_queries_logits = outputs.masks_queries_logits
    masks_queries_logits = torch.nn.functional.interpolate(masks_queries_logits, size=(384, 384), mode='bilinear', align_corners=False)
    masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1]
    masks_probs = masks_queries_logits.sigmoid()
    segmentation = torch.einsum('bqc, bqhw -> bchw', masks_classes, masks_probs)
    batch_size = class_queries_logits.shape[0]
    if target_sizes is not None:
        if batch_size != len(target_sizes):
            raise ValueError('Make sure that you pass in as many target sizes as the batch dimension of the logits')
        semantic_segmentation = []
        for idx in range(batch_size):
            resized_logits = torch.nn.functional.interpolate(segmentation[idx].unsqueeze(dim=0), size=target_sizes[idx], mode='bilinear', align_corners=False)
            semantic_map = resized_logits[0].argmax(dim=0)
            semantic_segmentation.append(semantic_map)
    else:
        semantic_segmentation = segmentation.argmax(dim=1)
        semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]
    return semantic_segmentation

Converts the output of [`Mask2FormerForUniversalSegmentation`] into semantic segmentation maps. Only supports
PyTorch.

Args:
    outputs ([`Mask2FormerForUniversalSegmentation`]):
        Raw outputs of the model.
    target_sizes (`List[Tuple[int, int]]`, *optional*):
        List of length (batch_size), where each list item (`Tuple[int, int]]`) corresponds to the requested
        final size (height, width) of each prediction. If left to None, predictions will not be resized.
Returns:
    `List[torch.Tensor]`:
        A list of length `batch_size`, where each item is a semantic segmentation map of shape (height, width)
        corresponding to the target_sizes entry (if `target_sizes` is specified). Each entry of each
        `torch.Tensor` correspond to a semantic class id.


def get_dataclass(self, json_dataclass: type[T]) -> T:
    if not mime_types.is_dataclass(self.mimetype):
        raise ValueError('Part is not a dataclass.')
    try:
        return json_dataclass.from_json(self.text)
    except AttributeError as e:
        raise ValueError(f'{json_dataclass.__name__} is not a valid json dataclass') from e

Returns representation of the Part as a given dataclass.

Args:
  json_dataclass: A dataclass that can be converted to/from JSON.

Returns:
  The dataclass representation of the Part.


def map_structure(func, *structure, **kwargs):
    return nest_util.map_structure(nest_util.Modality.CORE, func, *structure, **kwargs)

Creates a new structure by applying `func` to each atom in `structure`.

Refer to [tf.nest](https://www.tensorflow.org/api_docs/python/tf/nest)
for the definition of a structure.

Applies `func(x[0], x[1], ...)` where x[i] enumerates all atoms in
`structure[i]`.  All items in `structure` must have the same arity,
and the return value will contain results with the same structure layout.

Examples:

* A single Python dict:

>>> a = {"hello": 24, "world": 76}
>>> tf.nest.map_structure(lambda p: p * 2, a)
{'hello': 48, 'world': 152}

* Multiple Python dictionaries:

>>> d1 = {"hello": 24, "world": 76}
>>> d2 = {"hello": 36, "world": 14}
>>> tf.nest.map_structure(lambda p1, p2: p1 + p2, d1, d2)
{'hello': 60, 'world': 90}

* A single Python list:

>>> a = [24, 76, "ab"]
>>> tf.nest.map_structure(lambda p: p * 2, a)
[48, 152, 'abab']

* Scalars:

>>> tf.nest.map_structure(lambda x, y: x + y, 3, 4)
7

* Empty structures:

>>> tf.nest.map_structure(lambda x: x + 1, ())
()

* Check the types of iterables:

>>> s1 = (((1, 2), 3), 4, (5, 6))
>>> s1_list = [[[1, 2], 3], 4, [5, 6]]
>>> tf.nest.map_structure(lambda x, y: None, s1, s1_list)
Traceback (most recent call last):
...
TypeError: The two structures don't have the same nested structure

* Type check is set to False:

>>> s1 = (((1, 2), 3), 4, (5, 6))
>>> s1_list = [[[1, 2], 3], 4, [5, 6]]
>>> tf.nest.map_structure(lambda x, y: None, s1, s1_list, check_types=False)
(((None, None), None), None, (None, None))

Args:
  func: A callable that accepts as many arguments as there are structures.
  *structure: atom or nested structure.
  **kwargs: Valid keyword args are:
    * `check_types`: If set to `True` (default) the types of iterables within
      the structures have to be same (e.g. `map_structure(func, [1], (1,))`
      raises a `TypeError` exception). To allow this set this argument to
      `False`. Note that namedtuples with identical name and fields are always
      considered to have the same shallow structure.
    * `expand_composites`: If set to `True`, then composite tensors such as
      `tf.sparse.SparseTensor` and `tf.RaggedTensor` are expanded into their
      component tensors.  If `False` (the default), then composite tensors are
      not expanded.

Returns:
  A new structure with the same arity as `structure[0]`, whose atoms
  correspond to `func(x[0], x[1], ...)` where `x[i]` is the atom in the
  corresponding location in `structure[i]`. If there are different structure
  types and `check_types` is `False` the structure types of the first
  structure will be used.

Raises:
  TypeError: If `func` is not callable or if the structures do not match
    each other by depth tree.
  ValueError: If no structure is provided or if the structures do not match
    each other by type.
  ValueError: If wrong keyword arguments are provided.


def build_graph(device, input_shape, output_sizes, axis):
    with ops.device('/%s:0' % device):
        inp = array_ops.zeros(input_shape)
        outputs = []
        for _ in range(100):
            outputs.extend(array_ops.split(inp, output_sizes, axis))
        return control_flow_ops.group(*outputs)

Build a graph containing a sequence of split operations.

Args:
  device: string, the device to run on.
  input_shape: shape of the input tensor.
  output_sizes: size of each output along axis.
  axis: axis to be split along.

Returns:
  An array of tensors to run()


def _AssertProtoEquals(self, a, b, msg=None, relative_tolerance=None):
    if not compare.ProtoEq(a, b):
        compare.assertProtoEqual(self, a, b, normalize_numbers=True, msg=msg, relative_tolerance=relative_tolerance)

Asserts that a and b are the same proto.

Uses ProtoEq() first, as it returns correct results
for floating point attributes, and then use assertProtoEqual()
in case of failure as it provides good error messages.

Args:
  a: a proto.
  b: another proto.
  msg: Optional message to report on failure.
  relative_tolerance: float. The allowable difference between the two values
    being compared is determined by multiplying the relative tolerance by
    the maximum of the two values. If this is not provided, then all floats
    are compared using string comparison.


def pmap(f, axis_name=None, devices=None):
    if devices is None:
        devices = accelerators()
    if not isinstance(devices, (list, tuple)):
        raise ValueError('Must pass a list or tuple of devices')
    num_devices = len(devices)
    if not num_devices:
        raise ValueError('There must be at least 1 device')
    has_tpu = bool(tpu_devices(devices))
    pmap_fn = _get_pmap_impl(f, devices, has_tpu)

    def wrapper(*args):
        """Wrapper that wraps/unwraps args, retvals, and runs the function."""
        if _pmap_config.devices() is not None:
            raise ValueError('Found a surrounding pmap. Nested pmap is not supported yet.')
        flattened_input_args = nest.flatten(args)
        flattened_per_device_args = [[] for _ in devices]
        for arg in flattened_input_args:
            if isinstance(arg, tensor_lib.Tensor):
                if not arg.shape.rank or arg.shape[0] != len(devices):
                    raise ValueError('Input tensors need to have a first dimension equal to the number of devices; got tensor of shape %s and %s devices' % (arg.shape, len(devices)))
                for j, device in enumerate(devices):
                    updated_arg = array_ops.gather_v2(arg, j)
                    if not has_tpu:
                        with ops.device(device):
                            updated_arg = array_ops.identity(updated_arg)
                    flattened_per_device_args[j].append(updated_arg)
            elif isinstance(arg, ShardedNdArray):
                for device_args, tensor in zip(flattened_per_device_args, arg.tensors):
                    device_args.append(tensor)
            else:
                for device_args in flattened_per_device_args:
                    device_args.append(arg)
        all_per_device_args = [nest.pack_sequence_as(args, device_args) for device_args in flattened_per_device_args]
        with pmap_config(axis_name, devices):
            results = pmap_fn(all_per_device_args)
        flattened_results = [nest.flatten(result) for result in results]
        final_tree = []
        for i in range(len(flattened_results[0])):
            tensors = []
            for j, device in enumerate(devices):
                assert isinstance(flattened_results[j][i], tensor_lib.Tensor), 'currently only tensor return items are supported'
                tensors.append(flattened_results[j][i])
            final_tree.append(ShardedNdArray(tensors))
        return nest.pack_sequence_as(results[0], final_tree)
    return wrapper

Transforms a function into a multi-device function.

The semantics are similar to JAX's pmap.

Args:
  f: The function to be converted.
  axis_name: Used for nested pmap, which is not supported yet.
  devices: The devices over which the returned function will run.

Returns:
  A function that runs the underlying function `f` on `devices`. Its arguments
  can be `ShardedNdArray`s, tensors or other Python objects, and its return
  values are all `ShardedNdArray`s. If an input is a tensor, the length of its
  first dimension must equal the number of devices, and the tensor will be
  splitted along its first dimension among the devices. If an input is an
  unknown Python object, it will be replicated among the devices.


def use_value_spec(self, value_spec: Optional[pg_typing.Dict], allow_partial: bool=False) -> 'Dict':
    if value_spec is None:
        self._value_spec = None
        self._accessor_writable = True
        return self
    if not isinstance(value_spec, pg_typing.Dict):
        raise ValueError(self._error_message(f'Value spec for list must be a `pg.typing.Dict` object. Encountered: {value_spec!r}'))
    if self._value_spec and self._value_spec != value_spec:
        raise RuntimeError(self._error_message(f'Dict is already bound with a different value spec: {self._value_spec}. New value spec: {value_spec}.'))
    self._allow_partial = allow_partial
    if flags.is_type_check_enabled():
        value_spec.apply(self, allow_partial=base.accepts_partial(self), child_transform=base.symbolic_transform_fn(self._allow_partial), root_path=self.sym_path)
    else:
        self._value_spec = value_spec
    return self

Applies a ``pg.typing.Dict`` as the value spec for current dict.

Args:
  value_spec: A Dict ValueSpec to apply to this Dict.
    If current Dict is schema-less (whose immediate members are not
    validated against schema), and `value_spec` is not None, the value spec
    will be applied to the Dict.
    Or else if current Dict is already symbolic (whose immediate members
    are under the constraint of a Dict value spec), and `value_spec` is
    None, current Dict will become schema-less. However, the schema
    constraints for non-immediate members will remain.
  allow_partial: Whether allow partial dict based on the schema. This flag
    will override allow_partial flag in __init__ for spec-less Dict.

Returns:
  Self.

Raises:
  ValueError: validation failed due to value error.
  RuntimeError: Dict is already bound with another spec.
  TypeError: type errors during validation.
  KeyError: key errors during validation.


class InputLayer(base_layer.Layer):

    def __init__(self, input_shape=None, batch_size=None, dtype=None, input_tensor=None, sparse=None, name=None, ragged=None, type_spec=None, **kwargs):
        self._init_input_shape = input_shape
        self._init_batch_size = batch_size
        self._init_dtype = dtype
        self._init_sparse = sparse
        self._init_ragged = ragged
        self._init_type_spec = type_spec
        strategy = distribute_lib.get_strategy()
        if strategy and batch_size is not None and distributed_training_utils.global_batch_size_supported(strategy):
            if batch_size % strategy.num_replicas_in_sync != 0:
                raise ValueError('The `batch_size` argument ({}) must be divisible by the number of replicas ({})'.format(batch_size, strategy.num_replicas_in_sync))
            batch_size = batch_size // strategy.num_replicas_in_sync
        if 'batch_input_shape' in kwargs:
            batch_input_shape = kwargs.pop('batch_input_shape')
            if input_shape and batch_input_shape:
                raise ValueError('Only provide the input_shape OR batch_input_shape argument to InputLayer, not both at the same time.')
            if batch_input_shape:
                batch_size = batch_input_shape[0]
                input_shape = batch_input_shape[1:]
        if kwargs:
            raise ValueError('Unrecognized keyword arguments:', kwargs.keys())
        if sparse and ragged:
            raise ValueError('Cannot set both sparse and ragged to True in a Keras input.')
        if not name:
            prefix = 'input'
            name = prefix + '_' + str(backend.get_uid(prefix))
        if not dtype:
            if input_tensor is None:
                dtype = backend.floatx()
            else:
                dtype = backend.dtype(input_tensor)
        elif input_tensor is not None and input_tensor.dtype != dtype:
            raise ValueError('`input_tensor.dtype` differs from `dtype`: %s vs. %s' % (input_tensor.dtype, dtype))
        super(InputLayer, self).__init__(dtype=dtype, name=name)
        self.built = True
        self.sparse = True if sparse else False
        self.ragged = True if ragged else False
        self.batch_size = batch_size
        self.supports_masking = True
        if isinstance(input_shape, tensor_shape.TensorShape):
            input_shape = tuple(input_shape.as_list())
        elif isinstance(input_shape, int):
            input_shape = (input_shape,)
        if type_spec is not None:
            args_that_must_be_none = [('(input_)shape', self._init_input_shape), ('batch_size', self._init_batch_size), ('dtype', self._init_dtype), ('input_tensor', input_tensor), ('sparse', self._init_sparse), ('ragged', self._init_ragged)]
            for arg_name, arg in args_that_must_be_none:
                _assert_other_arg_none(arg_name, arg)
            if not ops.executing_eagerly_outside_functions():
                raise ValueError('Creating Keras inputs from a type_spec is only supported when eager execution is enabled.')
            input_tensor = keras_tensor.keras_tensor_from_type_spec(type_spec)
            if isinstance(input_tensor, keras_tensor.SparseKerasTensor):
                self.sparse = True
            if isinstance(input_tensor, keras_tensor.RaggedKerasTensor):
                self.ragged = True
            self.is_placeholder = True
            try:
                self._batch_input_shape = tuple(input_tensor.shape.as_list())
            except ValueError:
                self._batch_input_shape = None
        elif input_tensor is None:
            if input_shape is not None:
                batch_input_shape = (batch_size,) + tuple(input_shape)
            else:
                batch_input_shape = None
            graph = backend.get_graph()
            with graph.as_default():
                input_tensor = backend.placeholder(shape=batch_input_shape, dtype=dtype, name=self.name, sparse=sparse, ragged=ragged)
            self.is_placeholder = True
            self._batch_input_shape = batch_input_shape
        else:
            if ops.executing_eagerly_outside_functions():
                if not isinstance(input_tensor, keras_tensor.KerasTensor):
                    input_tensor = keras_tensor.keras_tensor_from_tensor(input_tensor)
            elif not tf_utils.is_symbolic_tensor(input_tensor):
                raise ValueError('You should not pass an EagerTensor to `Input`. For example, instead of creating an InputLayer, you should instantiate your model and directly call it on your input.')
            self.is_placeholder = False
            try:
                self._batch_input_shape = tuple(input_tensor.shape.as_list())
            except ValueError:
                self._batch_input_shape = None
        input_tensor._keras_mask = None
        node_module.Node(layer=self, outputs=input_tensor)
        if isinstance(input_tensor, keras_tensor.KerasTensor) or tf_utils.is_extension_type(input_tensor):
            self._type_spec = input_tensor._type_spec
        else:
            self._type_spec = tensor_spec.TensorSpec(shape=input_tensor.shape, dtype=input_tensor.dtype, name=self.name)

    def get_config(self):
        if self._init_type_spec is not None:
            config = {'name': self.name, 'type_spec': self._init_type_spec}
        else:
            config = {'batch_input_shape': self._batch_input_shape, 'dtype': self.dtype, 'sparse': self.sparse, 'ragged': self.ragged, 'name': self.name}
        return config

    @property
    def _trackable_saved_model_saver(self):
        return layer_serialization.InputLayerSavedModelSaver(self)

Layer to be used as an entry point into a Network (a graph of layers).

It can either wrap an existing tensor (pass an `input_tensor` argument)
or create a placeholder tensor (pass arguments `input_shape`, and
optionally, `dtype`).

It is generally recommend to use the functional layer API via `Input`,
(which creates an `InputLayer`) without directly using `InputLayer`.

When using InputLayer with Keras Sequential model, it can be skipped by
moving the input_shape parameter to the first layer after the InputLayer.

This class can create placeholders for tf.Tensors, tf.SparseTensors, and
tf.RaggedTensors by choosing 'sparse=True' or 'ragged=True'. Note that
'sparse' and 'ragged' can't be configured to True at same time.
Usage:

```python
# With explicit InputLayer.
model = tf.keras.Sequential([
  tf.keras.layers.InputLayer(input_shape=(4,)),
  tf.keras.layers.Dense(8)])
model.compile(tf.optimizers.RMSprop(0.001), loss='mse')
model.fit(np.zeros((10, 4)),
          np.ones((10, 8)))

# Without InputLayer and let the first layer to have the input_shape.
# Keras will add a input for the model behind the scene.
model = tf.keras.Sequential([
  tf.keras.layers.Dense(8, input_shape=(4,))])
model.compile(tf.optimizers.RMSprop(0.001), loss='mse')
model.fit(np.zeros((10, 4)),
          np.ones((10, 8)))
```

Args:
    input_shape: Shape tuple (not including the batch axis), or `TensorShape`
      instance (not including the batch axis).
    batch_size: Optional input batch size (integer or None).
    dtype: Optional datatype of the input. When not provided, the Keras
        default float type will be used.
    input_tensor: Optional tensor to use as layer input. If set, the layer
        will use the `tf.TypeSpec` of this tensor rather
        than creating a new placeholder tensor.
    sparse: Boolean, whether the placeholder created is meant to be sparse.
        Default to False.
    ragged: Boolean, whether the placeholder created is meant to be ragged.
        In this case, values of 'None' in the 'shape' argument represent
        ragged dimensions. For more information about RaggedTensors, see
        [this guide](https://www.tensorflow.org/guide/ragged_tensors).
        Default to False.
    type_spec: A `tf.TypeSpec` object to create Input from. This `tf.TypeSpec`
        represents the entire batch. When provided, all other args except
        name must be None.
    name: Optional name of the layer (string).


class DType(abc.ABC):
    name: str
    array_cls_name: str

    @classmethod
    def from_value(cls, value: Any) -> DType:
        """Convert the value to dtype."""
        if value is None:
            return AnyDType()
        elif isinstance(value, DType):
            return value
        elif numpy_utils.lazy.is_dtype(value):
            return NpDType(numpy_utils.lazy.as_dtype(value))
        elif value in _STD_TYPE_TO_DTYPE:
            return _STD_TYPE_TO_DTYPE[value]
        else:
            raise TypeError(f'Unsuported dtype: {value!r}')

    def asarray(self, array_like, *, xnp: numpy_utils.NpModule, casting: Union[Casting, str]=Casting.ALL):
        """Creates an `xnp.ndarray` from the `array_like`.

    Args:
      array_like: Any array-like
      xnp: Target numpy module
      casting: If `NONE`, prevent casting.

    Returns:
      array: The xnp array.
    """
        casting = Casting(casting)
        from_dtype = numpy_utils.lazy.dtype_from_array(array_like, strict=False)
        to_dtype = self._get_target_dtype(from_dtype)
        if casting == casting.NONE:
            if to_dtype is None:
                pass
            elif from_dtype is None:
                pass
            elif from_dtype != to_dtype:
                raise ValueError(f'Cannot cast {from_dtype} to {to_dtype} (casting={casting}).')
        elif casting == casting.ALL:
            pass
        else:
            raise NotImplementedError(f'Unsupported casting {casting}')
        if to_dtype is None:
            dtype_kwargs = {}
        else:
            dtype_kwargs = {'dtype': numpy_utils.lazy.as_dtype(to_dtype, xnp=xnp)}
        if isinstance(array_like, np.ndarray) and array_like.shape == ():
            if not dtype_kwargs:
                dtype_kwargs = {'dtype': numpy_utils.lazy.as_dtype(array_like.dtype, xnp=xnp)}
            array_like = array_like.item()
        arr = xnp.asarray(array_like, **dtype_kwargs)
        return arr

    @abc.abstractmethod
    def _get_target_dtype(self, from_dtype: Optional[_NpDType]) -> Optional[_NpDType]:
        """Validate and normalize the numpy dtype.

    Args:
      from_dtype: DType of the array to cast

    Returns:
      to_dtype: DType of the array after casting
    """

    @abc.abstractmethod
    def __eq__(self, other) -> bool:
        raise NotImplementedError

    @abc.abstractmethod
    def __hash__(self) -> int:
        raise NotImplementedError

    def __repr__(self) -> str:
        return f'DType({self.name})'

DType wrapper.

This allow to support more complex types, like dtype unions.

This is EXPERIMENTAL, so the API might change.

Attributes:
  name: Representation name (e.g. np.uint8, AnyFloat...)
  array_cls_name: Name of the array class associated with the dtype (`f32`,
    `ui8`,...).


Called at the end of evaluation or validation.

Subclasses should override for any actions to run.

Args:
    logs: Dict. Currently the output of the last call to
      `on_test_batch_end()` is passed to this argument for this method
      but that may change in the future.


def train_step(self, data):
    data = data_adapter.expand_1d(data)
    x, y, sample_weight = data_adapter.unpack_x_y_sample_weight(data)
    with backprop.GradientTape() as tape:
        y_pred = self(x, training=True)
        loss = self.compiled_loss(y, y_pred, sample_weight, regularization_losses=self.losses)
    self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
    self.compiled_metrics.update_state(y, y_pred, sample_weight)
    return_metrics = {}
    for metric in self.metrics:
        result = metric.result()
        if isinstance(result, dict):
            return_metrics.update(result)
        else:
            return_metrics[metric.name] = result
    return return_metrics

The logic for one training step.

This method can be overridden to support custom training logic.
For concrete examples of how to override this method see
[Customizing what happens in
fit](https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit).
This method is called by `Model.make_train_function`.

This method should contain the mathematical logic for one step of training.
This typically includes the forward pass, loss calculation, backpropagation,
and metric updates.

Configuration details for *how* this logic is run (e.g. `tf.function` and
`tf.distribute.Strategy` settings), should be left to
`Model.make_train_function`, which can also be overridden.

Args:
  data: A nested structure of `Tensor`s.

Returns:
  A `dict` containing values that will be passed to
  `tf.keras.callbacks.CallbackList.on_train_batch_end`. Typically, the
  values of the `Model`'s metrics are returned. Example:
  `{'loss': 0.2, 'accuracy': 0.7}`.


def __call__(self, **kwargs):
    if len(kwargs) != len(self._inputs):
        raise ValueError('Invalid number of inputs provided for running a SignatureDef, expected %s vs provided %s' % (len(self._inputs), len(kwargs)))
    for input_name, value in kwargs.items():
        if input_name not in self._inputs:
            raise ValueError('Invalid Input name (%s) for SignatureDef' % input_name)
        self._interpreter_wrapper.ResizeInputTensor(self._inputs[input_name], np.array(value.shape, dtype=np.int32), False, self._subgraph_index)
    self._interpreter_wrapper.AllocateTensors(self._subgraph_index)
    for input_name, value in kwargs.items():
        self._interpreter_wrapper.SetTensor(self._inputs[input_name], value, self._subgraph_index)
    self._interpreter_wrapper.Invoke(self._subgraph_index)
    result = {}
    for output_name, output_index in self._outputs:
        result[output_name] = self._interpreter_wrapper.GetTensor(output_index, self._subgraph_index)
    return result

Runs the SignatureDef given the provided inputs in arguments.

Args:
  **kwargs: key,value for inputs to the model. Key is the SignatureDef input
    name. Value is numpy array with the value.

Returns:
  dictionary of the results from the model invoke.
  Key in the dictionary is SignatureDef output name.
  Value is the result Tensor.


def read_video_opencv(video_path: str, sample_indices_fn: Callable, **kwargs):
    requires_backends(read_video_opencv, ['cv2'])
    import cv2
    video = cv2.VideoCapture(video_path)
    total_num_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    video_fps = video.get(cv2.CAP_PROP_FPS)
    duration = total_num_frames / video_fps if video_fps else 0
    metadata = VideoMetadata(total_num_frames=int(total_num_frames), fps=float(video_fps), duration=float(duration), video_backend='opencv')
    indices = sample_indices_fn(metadata=metadata, **kwargs)
    index = 0
    frames = []
    while video.isOpened():
        success, frame = video.read()
        if not success:
            break
        if index in indices:
            height, width, channel = frame.shape
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frames.append(frame[0:height, 0:width, 0:channel])
        if success:
            index += 1
        if index >= total_num_frames:
            break
    video.release()
    metadata.frames_indices = indices
    return (np.stack(frames), metadata)

Decode a video using the OpenCV backend.

Args:
    video_path (`str`):
        Path to the video file.
    sample_indices_fn (`Callable`):
        A callable function that will return indices at which the video should be sampled. If the video has to be loaded using
        by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`.
        If not provided, simple uniform sampling with fps is performed.
        Example:
        def sample_indices_fn(metadata, **kwargs):
            return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int)

Returns:
    Tuple[`np.array`, `VideoMetadata`]: A tuple containing:
        - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]).
        - `VideoMetadata` object.


def __init__(self, columns: list[str], min_value: float=0.0, max_value: float=1.0, name: Optional[str]=None):
    super().__init__(columns)
    self.min_value = min_value
    self.max_value = max_value
    self.name = name
    if self.max_value <= self.min_value:
        raise ValueError('max_value must be greater than min_value')

This function applies a scaling transformation on the given columns
of incoming data. The transformation scales the input values to the
range [min_value, max_value].

Args:
  columns: A list of column names to apply the transformation on.
  min_value: The minimum value of the output range.
  max_value: The maximum value of the output range.
  name: A name for the operation (optional).


def _prepare_4d_causal_attention_mask_with_cache_position(attention_mask: torch.Tensor, sequence_length: int, target_length: int, dtype: torch.dtype, cache_position: torch.Tensor, batch_size: int, **kwargs):
    if attention_mask is not None and attention_mask.dim() == 4:
        causal_mask = attention_mask
    else:
        min_dtype = torch.finfo(dtype).min
        causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device)
        if sequence_length != 1:
            causal_mask = torch.triu(causal_mask, diagonal=1)
        causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1)
        causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
        if attention_mask is not None:
            causal_mask = causal_mask.clone()
            mask_length = attention_mask.shape[-1]
            padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(causal_mask.device)
            padding_mask = padding_mask == 0
            causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(padding_mask, min_dtype)
    return causal_mask

Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
`(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.

Args:
    attention_mask (`torch.Tensor`):
        A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape
        `(batch_size, 1, query_length, key_value_length)`.
    sequence_length (`int`):
        The sequence length being processed.
    target_length (`int`):
        The target length: when generating with static cache, the mask should be as long as the static cache,
        to account for the 0 padding, the part of the cache that is not filled yet.
    dtype (`torch.dtype`):
        The dtype to use for the 4D attention mask.
    cache_position (`torch.Tensor`):
        Indices depicting the position of the input sequence tokens in the sequence.
    batch_size (`torch.Tensor`):
        Batch size.


def __call__(self, raw_speech: Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]], padding: Union[bool, str, PaddingStrategy]=True, pad_to_multiple_of: Optional[int]=2, max_length: Optional[int]=None, truncation: bool=False, return_tensors: Optional[Union[str, TensorType]]=None, sampling_rate: Optional[int]=None, return_attention_mask: Optional[bool]=None, do_normalize_per_mel_bins: Optional[bool]=True, **kwargs) -> BatchFeature:
    if sampling_rate is not None:
        if sampling_rate != self.sampling_rate:
            raise ValueError(f'The model corresponding to this feature extractor: {self} was trained using a sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled with {self.sampling_rate} and not {sampling_rate}.')
    else:
        logger.warning(f'It is strongly recommended to pass the `sampling_rate` argument to `{self.__class__.__name__}()`. Failing to do so can result in silent errors that might be hard to debug.')
    return_attention_mask = return_attention_mask if return_attention_mask is not None else self.return_attention_mask
    is_batched_numpy = isinstance(raw_speech, np.ndarray) and len(raw_speech.shape) > 1
    if is_batched_numpy and len(raw_speech.shape) > 3:
        raise ValueError(f'Only mono-channel or stereo-channel audio is supported for input to {self}')
    acceptable_types = (torch.Tensor, np.ndarray, tuple, list) if is_torch_available() else (np.ndarray, tuple, list)
    is_batched = is_batched_numpy or (isinstance(raw_speech, (list, tuple)) and isinstance(raw_speech[0], acceptable_types))
    if is_batched:
        raw_speech = [np.asarray(speech, dtype=np.float32) for speech in raw_speech]
    elif not is_batched and (not isinstance(raw_speech, np.ndarray)):
        raw_speech = np.asarray(raw_speech, dtype=np.float32)
    elif isinstance(raw_speech, np.ndarray) and raw_speech.dtype is np.dtype(np.float64):
        raw_speech = raw_speech.astype(np.float32)
    if not is_batched:
        raw_speech = [raw_speech]
    features = [self._extract_fbank_features(waveform) for waveform in raw_speech]
    if do_normalize_per_mel_bins:
        features = [(x - np.expand_dims(x.mean(0), 0)) / np.sqrt(np.expand_dims(x.var(0, ddof=1), 0) + 1e-07) for x in features]
    encoded_inputs = BatchFeature({'input_features': features})
    padded_inputs = self.pad(encoded_inputs, padding=padding, max_length=max_length, truncation=truncation, pad_to_multiple_of=pad_to_multiple_of, return_attention_mask=True, return_tensors='np')
    input_features = padded_inputs.get('input_features')
    attention_mask = padded_inputs.pop('attention_mask')
    batch_size, num_frames, num_channels = input_features.shape
    remainder = num_frames % self.stride
    if remainder != 0:
        input_features = input_features[:, :num_frames - remainder, :]
        attention_mask = attention_mask[:, :num_frames - remainder]
    input_features = np.reshape(input_features, (batch_size, num_frames // self.stride, num_channels * self.stride))
    indices = np.arange(0, num_frames - remainder)
    attention_mask = attention_mask[:, indices % self.stride == 1]
    padded_inputs['input_features'] = input_features
    if return_attention_mask:
        padded_inputs['attention_mask'] = attention_mask
    if return_tensors is not None:
        padded_inputs = padded_inputs.convert_to_tensors(return_tensors)
    return padded_inputs

Main method to featurize and prepare for the model one or several sequence(s).

Args:
    raw_speech (`np.ndarray`, `torch.Tensor`, `List[float]`, `List[np.ndarray]`, `List[torch.Tensor]`,
    `List[List[float]]`, `List[List[List[float]]]`):
        The sequence or batch of sequences to be padded. Each sequence can be a numpy array,
        a torch tensor, a list of float values, a list of numpy arrays, a list of torch tensors,
        a list of list of float values or a list of a list of list of float values.
        If `raw_speech` is a one-dimensional `np.ndarray`, `torch.Tensor` or a `List[float]`, `raw_speech` is
        considered a single-channel, single-sample sound. In all other cases, the first dimension of
        `raw_speech`, whether from an `np.ndarray`, a `torch.Tensor` or a `List[...]`,
        corresponds to the number of samples in the batch, and the number of channels
        (i.e. mono or stereo character) is derived from the other dimensions
        (1D -> single-channel waveform batches; 2D-> stereo-channel waveform batches).
    padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):
        Select a strategy to pad the returned sequences (according to the model's padding side and padding
        index) among:

        - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
          sequence if provided).
        - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
          acceptable input length for the model if that argument is not provided.
        - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
          lengths).
    pad_to_multiple_of (`int`, *optional*, defaults to 2):
        If set will pad the sequence to a multiple of the provided value.

        This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability
        `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128.
    max_length (`int`, *optional*):
        Maximum length of the returned list and optionally padding length (see above).
    truncation (`bool`):
        Activates truncation to cut input sequences longer than *max_length* to *max_length*.
    return_attention_mask (`bool`, *optional*):
        Whether to return the attention mask. If left to the default, will return the attention mask according
        to the specific feature_extractor's default.

        [What are attention masks?](../glossary#attention-mask)

        <Tip>

        For SeamlessM4T models, `attention_mask` should always be passed for batched inference, to avoid subtle
        bugs.

        </Tip>

    return_tensors (`str` or [`~utils.TensorType`], *optional*):
        If set, will return tensors instead of list of python integers. Acceptable values are:

        - `'tf'`: Return TensorFlow `tf.constant` objects.
        - `'pt'`: Return PyTorch `torch.Tensor` objects.
        - `'np'`: Return Numpy `np.ndarray` objects.
    sampling_rate (`int`, *optional*):
        The sampling rate at which the `raw_speech` input was sampled. It is strongly recommended to pass
        `sampling_rate` at the forward call to prevent silent errors.
    do_normalize_per_mel_bins (`bool`, *optional*, defaults to `True`):
        Whether or not to zero-mean unit-variance normalize the input per mel-channel.
    kwargs (*optional*):
        Remaining dictionary of keyword arguments that will be passed to the tokenizer or the feature
        extractor.


def compute_iou(boxes1, boxes2, bounding_box_format, use_masking=False, mask_val=-1, image_shape=None):
    boxes1_rank = len(ops.shape(boxes1))
    boxes2_rank = len(ops.shape(boxes2))
    if boxes1_rank not in [2, 3]:
        raise ValueError(f'compute_iou() expects boxes1 to be batched, or to be unbatched. Received len(boxes1.shape)={boxes1_rank}, len(boxes2.shape)={boxes2_rank}. Expected either len(boxes1.shape)=2 AND or len(boxes1.shape)=3.')
    if boxes2_rank not in [2, 3]:
        raise ValueError(f'compute_iou() expects boxes2 to be batched, or to be unbatched. Received len(boxes1.shape)={boxes1_rank}, len(boxes2.shape)={boxes2_rank}. Expected either len(boxes2.shape)=2 AND or len(boxes2.shape)=3.')
    target_format = 'yxyx'
    if 'rel' in bounding_box_format and image_shape is None:
        raise ValueError(f'When using relative bounding box formats (e.g. `rel_yxyx`) the `image_shape` argument must be provided.Received `image_shape`: {image_shape}')
    if image_shape is None:
        height, width = (None, None)
    else:
        height, width, _ = image_shape
    boxes1 = converters.convert_format(boxes1, source=bounding_box_format, target=target_format, height=height, width=width)
    boxes2 = converters.convert_format(boxes2, source=bounding_box_format, target=target_format, height=height, width=width)
    intersect_area = _compute_intersection(boxes1, boxes2)
    boxes1_area = _compute_area(boxes1)
    boxes2_area = _compute_area(boxes2)
    boxes2_area_rank = len(boxes2_area.shape)
    boxes2_axis = 1 if boxes2_area_rank == 2 else 0
    boxes1_area = ops.expand_dims(boxes1_area, axis=-1)
    boxes2_area = ops.expand_dims(boxes2_area, axis=boxes2_axis)
    union_area = boxes1_area + boxes2_area - intersect_area
    res = ops.divide(intersect_area, union_area + backend.epsilon())
    if boxes1_rank == 2:
        perm = [1, 0]
    else:
        perm = [0, 2, 1]
    if not use_masking:
        return res
    mask_val_t = ops.cast(mask_val, res.dtype) * ops.ones_like(res)
    boxes1_mask = ops.less(ops.max(boxes1, axis=-1, keepdims=True), 0.0)
    boxes2_mask = ops.less(ops.max(boxes2, axis=-1, keepdims=True), 0.0)
    background_mask = ops.logical_or(boxes1_mask, ops.transpose(boxes2_mask, perm))
    iou_lookup_table = ops.where(background_mask, mask_val_t, res)
    return iou_lookup_table

Computes a lookup table vector containing the ious for a given set boxes.

The lookup vector is to be indexed by [`boxes1_index`,`boxes2_index`] if
boxes are unbatched and by [`batch`, `boxes1_index`,`boxes2_index`] if the
boxes are batched.

The users can pass `boxes1` and `boxes2` to be different ranks. For example:
1) `boxes1`: [batch_size, M, 4], `boxes2`: [batch_size, N, 4] -> return
    [batch_size, M, N].
2) `boxes1`: [batch_size, M, 4], `boxes2`: [N, 4] -> return
    [batch_size, M, N]
3) `boxes1`: [M, 4], `boxes2`: [batch_size, N, 4] -> return
    [batch_size, M, N]
4) `boxes1`: [M, 4], `boxes2`: [N, 4] -> return [M, N]

Args:
    boxes1: a list of bounding boxes in 'corners' format. Can be batched or
        unbatched.
    boxes2: a list of bounding boxes in 'corners' format. Can be batched or
        unbatched.
    bounding_box_format: a case-insensitive string which is one of `"xyxy"`,
        `"rel_xyxy"`, `"xyWH"`, `"center_xyWH"`, `"yxyx"`, `"rel_yxyx"`.
        For detailed information on the supported format, see the
    use_masking: whether masking will be applied. This will mask all
        `boxes1` or `boxes2` that have values less than 0 in all its 4
        dimensions. Default to `False`.
    mask_val: int to mask those returned IOUs if the masking is True,
        defaults to -1.
    image_shape: `Tuple[int]`. The shape of the image (height, width, 3).
        When using relative bounding box format for `box_format` the
        `image_shape` is used for normalization.

Returns:
    iou_lookup_table: a vector containing the pairwise ious of boxes1 and
        boxes2.


def equity_swap_price(rate_leg_coupon_rates, equity_leg_forward_prices, equity_leg_spots, rate_leg_notional, equity_leg_notional, rate_leg_daycount_fractions, rate_leg_discount_factors, equity_leg_discount_factors, equity_dividends=None, is_equity_receiver=None, dtype=None, name=None):
    name = name or 'equity_swap_price'
    with tf.name_scope(name):
        rate_leg_coupon_rates = tf.convert_to_tensor(rate_leg_coupon_rates, dtype=dtype, name='rate_leg_coupon_rates')
        dtype = dtype or rate_leg_coupon_rates.dtype
        equity_leg_forward_prices = tf.convert_to_tensor(equity_leg_forward_prices, dtype=dtype, name='equity_leg_forward_prices')
        equity_leg_spots = tf.convert_to_tensor(equity_leg_spots, dtype=dtype, name='equity_leg_spots')
        rate_leg_daycount_fractions = tf.convert_to_tensor(rate_leg_daycount_fractions, dtype=dtype, name='rate_leg_daycount_fractions')
        equity_dividends = equity_dividends or 0
        equity_dividends = tf.convert_to_tensor(equity_dividends, dtype=dtype, name='equity_dividends')
        rate_leg_notional = tf.convert_to_tensor(rate_leg_notional, dtype=dtype, name='rate_leg_notional')
        equity_leg_notional = tf.convert_to_tensor(equity_leg_notional, dtype=dtype, name='equity_leg_notional')
        rate_leg_discount_factors = tf.convert_to_tensor(rate_leg_discount_factors, dtype=dtype, name='rate_leg_discount_factors')
        equity_leg_discount_factors = tf.convert_to_tensor(equity_leg_discount_factors, dtype=dtype, name='equity_leg_discount_factors')
        if is_equity_receiver is None:
            is_equity_receiver = True
        is_equity_receiver = tf.convert_to_tensor(is_equity_receiver, dtype=tf.bool, name='is_equity_receiver')
        one = tf.ones([], dtype=dtype)
        equity_receiver = tf.where(is_equity_receiver, one, -one)
        equity_cashflows = equity_leg_cashflows(forward_prices=equity_leg_forward_prices, spots=equity_leg_spots, notional=equity_leg_notional, dividends=equity_dividends)
        rate_cashflows = rate_leg_cashflows(coupon_rates=rate_leg_coupon_rates, notional=rate_leg_notional, daycount_fractions=rate_leg_daycount_fractions)
        return equity_receiver * swap_price(rate_cashflows, equity_cashflows, rate_leg_discount_factors, equity_leg_discount_factors)

Computes prices of a batch of equity swaps.

The swap consists of an equity and interest rate legs.

#### Example
```python
rate_leg_coupon_rates = [[0.1, 0.2, 0.05], [0.1, 0.05, 0.2]]
# Two cashflows of 4 and 3 payments, respectively
forward_prices = [[110, 120, 140, 150], [210, 220, 240, 0]]
spots = [100, 200]
notional = 1000
pay_leg_daycount_fractions = 0.5
rate_leg_daycount_fractions = [[0.5, 0.5, 0.5], [0.4, 0.5, 0.6]]
rate_leg_discount_factors = [[0.95, 0.9, 0.85], [0.98, 0.92, 0.88]]
equity_leg_discount_factors = [[0.95, 0.9, 0.85, 0.8],
                               [0.98, 0.92, 0.88, 0.0]]

equity_swap_price(
    rate_leg_coupon_rates=rate_leg_coupon_rates,
    equity_leg_forward_prices=forward_prices,
    equity_leg_spots=spots,
    rate_leg_notional=notional,
    equity_leg_notional=notional,
    rate_leg_daycount_fractions=rate_leg_daycount_fractions,
    rate_leg_discount_factors=rate_leg_discount_factors,
    equity_leg_discount_factors=equity_leg_discount_factors,
    is_equity_receiver=[True, False],
    dtype=tf.float64)
# Expected: [216.87770563, -5.00952381]
forward_rates(df_start_dates, df_end_dates, daycount_fractions,
              dtype=tf.float64)
```

Args:
  rate_leg_coupon_rates: A real `Tensor` of shape
    `batch_shape + [num_rate_cashflows]`, where `num_rate_cashflows` is the
    number of cashflows for each batch element. Coupon rates for the
    interest rate leg.
  equity_leg_forward_prices: A `Tensor` of the same `dtype` as
    `rate_leg_coupon_rates` and of shape
    `batch_shape + [num_equity_cashflows]`, where `num_equity_cashflows` is
    the number of cashflows for each batch element. Equity leg forward
    prices.
  equity_leg_spots: A `Tensor` of the same `dtype` as
    `equity_leg_forward_prices` and of shape compatible with `batch_shape`.
    Spot prices for each batch element of the equity leg.
  rate_leg_notional: A `Tensor` of the same `dtype` as `rate_leg_coupon_rates`
    and of compatible shape. Notional amount for each cashflow.
  equity_leg_notional: A `Tensor` of the same `dtype` as
    `equity_leg_forward_prices` and of compatible shape.  Notional amount for
    each cashflow.
  rate_leg_daycount_fractions: A `Tensor` of the same `dtype` as
    `rate_leg_coupon_rates` and of compatible shape.  Year fractions for the
    coupon accrual.
  rate_leg_discount_factors: A `Tensor` of the same `dtype` as
    `rate_leg_coupon_rates` and of compatible shape. Discount factors for each
    cashflow of the rate leg.
  equity_leg_discount_factors: A `Tensor` of the same `dtype` as
    `equity_leg_forward_prices` and of compatible shape. Discount factors for
    each cashflow of the equity leg.
  equity_dividends: A `Tensor` of the same `dtype` as
    `equity_leg_forward_prices` and of compatible shape. Dividends paid at the
    leg reset times.
    Default value: None which maps to zero dividend.
  is_equity_receiver: A boolean `Tensor` of shape compatible with
    `batch_shape`. Indicates whether the swap holder is equity holder or
    receiver.
    Default value: None which means that all swaps are equity reiver swaps.
  dtype: `tf.Dtype`. If supplied the dtype for the input and output `Tensor`s.
    Default value: None which maps to the default dtype inferred from
    `rate_leg_coupon_rates`.
  name: Python str. The name to give to the ops created by this function.
    Default value: None which maps to 'equity_swap_price'.

Returns:
  A `Tensor` of the same `dtype` as `rate_leg_coupon_rates` and of shape
  `batch_shape`. Present values of the equity swaps.


def add_content_to_file(file_name: Union[str, os.PathLike], content: str, add_after: Optional[Union[str, Pattern]]=None, add_before: Optional[Union[str, Pattern]]=None, exact_match: bool=False):
    with open(file_name, 'r', encoding='utf-8') as f:
        old_content = f.read()
    new_content = add_content_to_text(old_content, content, add_after=add_after, add_before=add_before, exact_match=exact_match)
    with open(file_name, 'w', encoding='utf-8') as f:
        f.write(new_content)

A utility to add some content inside a given file.

Args:
   file_name (`str` or `os.PathLike`): The name of the file in which we want to insert some content.
   content (`str`): The content to add.
   add_after (`str` or `Pattern`):
       The pattern to test on a line of `text`, the new content is added after the first instance matching it.
   add_before (`str` or `Pattern`):
       The pattern to test on a line of `text`, the new content is added before the first instance matching it.
   exact_match (`bool`, *optional*, defaults to `False`):
       A line is considered a match with `add_after` or `add_before` if it matches exactly when `exact_match=True`,
       otherwise, if `add_after`/`add_before` is present in the line.

<Tip warning={true}>

The arguments `add_after` and `add_before` are mutually exclusive, and one exactly needs to be provided.

</Tip>


def random_flip_up_down(image, seed=None):
    random_func = functools.partial(random_ops.random_uniform, seed=seed)
    return _random_flip(image, 0, random_func, 'random_flip_up_down')

Randomly flips an image vertically (upside down).

With a 1 in 2 chance, outputs the contents of `image` flipped along the first
dimension, which is `height`.  Otherwise, output the image as-is.
When passing a batch of images, each image will be randomly flipped
independent of other images.

Example usage:

>>> image = np.array([[[1], [2]], [[3], [4]]])
>>> tf.image.random_flip_up_down(image, 3).numpy().tolist()
[[[3], [4]], [[1], [2]]]

Randomly flip multiple images.

>>> images = np.array(
... [
...     [[[1], [2]], [[3], [4]]],
...     [[[5], [6]], [[7], [8]]]
... ])
>>> tf.image.random_flip_up_down(images, 4).numpy().tolist()
[[[[3], [4]], [[1], [2]]], [[[5], [6]], [[7], [8]]]]

For producing deterministic results given a `seed` value, use
`tf.image.stateless_random_flip_up_down`. Unlike using the `seed` param
with `tf.image.random_*` ops, `tf.image.stateless_random_*` ops guarantee the
same results given the same seed independent of how many times the function is
called, and independent of global seed settings (e.g. tf.random.set_seed).

Args:
  image: 4-D Tensor of shape `[batch, height, width, channels]` or 3-D Tensor
    of shape `[height, width, channels]`.
  seed: A Python integer. Used to create a random seed. See
    `tf.compat.v1.set_random_seed` for behavior.

Returns:
  A tensor of the same type and shape as `image`.
Raises:
  ValueError: if the shape of `image` not supported.


class AdamWeightDecay(Adam):

    def __init__(self, learning_rate: Union[float, schedules.LearningRateSchedule]=0.001, beta_1: float=0.9, beta_2: float=0.999, epsilon: float=1e-07, amsgrad: bool=False, weight_decay_rate: float=0.0, include_in_weight_decay: Optional[list[str]]=None, exclude_from_weight_decay: Optional[list[str]]=None, name: str='AdamWeightDecay', **kwargs):
        super().__init__(learning_rate, beta_1, beta_2, epsilon, amsgrad, name, **kwargs)
        self.weight_decay_rate = weight_decay_rate
        self._include_in_weight_decay = include_in_weight_decay
        self._exclude_from_weight_decay = exclude_from_weight_decay

    @classmethod
    def from_config(cls, config):
        """Creates an optimizer from its config with WarmUp custom object."""
        custom_objects = {'WarmUp': WarmUp}
        return super().from_config(config, custom_objects=custom_objects)

    def _prepare_local(self, var_device, var_dtype, apply_state):
        super()._prepare_local(var_device, var_dtype, apply_state)
        apply_state[var_device, var_dtype]['weight_decay_rate'] = tf.constant(self.weight_decay_rate, name='adam_weight_decay_rate')

    def _decay_weights_op(self, var, learning_rate, apply_state):
        do_decay = self._do_use_weight_decay(var.name)
        if do_decay:
            return var.assign_sub(learning_rate * var * apply_state[var.device, var.dtype.base_dtype]['weight_decay_rate'], use_locking=self._use_locking)
        return tf.no_op()

    def apply_gradients(self, grads_and_vars, name=None, **kwargs):
        grads, tvars = list(zip(*grads_and_vars))
        return super().apply_gradients(zip(grads, tvars), name=name, **kwargs)

    def _get_lr(self, var_device, var_dtype, apply_state):
        """Retrieves the learning rate with the given state."""
        if apply_state is None:
            return (self._decayed_lr_t[var_dtype], {})
        apply_state = apply_state or {}
        coefficients = apply_state.get((var_device, var_dtype))
        if coefficients is None:
            coefficients = self._fallback_apply_state(var_device, var_dtype)
            apply_state[var_device, var_dtype] = coefficients
        return (coefficients['lr_t'], {'apply_state': apply_state})

    def _resource_apply_dense(self, grad, var, apply_state=None):
        lr_t, kwargs = self._get_lr(var.device, var.dtype.base_dtype, apply_state)
        decay = self._decay_weights_op(var, lr_t, apply_state)
        with tf.control_dependencies([decay]):
            return super()._resource_apply_dense(grad, var, **kwargs)

    def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
        lr_t, kwargs = self._get_lr(var.device, var.dtype.base_dtype, apply_state)
        decay = self._decay_weights_op(var, lr_t, apply_state)
        with tf.control_dependencies([decay]):
            return super()._resource_apply_sparse(grad, var, indices, **kwargs)

    def get_config(self):
        config = super().get_config()
        config.update({'weight_decay_rate': self.weight_decay_rate})
        return config

    def _do_use_weight_decay(self, param_name):
        """Whether to use L2 weight decay for `param_name`."""
        if self.weight_decay_rate == 0:
            return False
        if self._include_in_weight_decay:
            for r in self._include_in_weight_decay:
                if re.search(r, param_name) is not None:
                    return True
        if self._exclude_from_weight_decay:
            for r in self._exclude_from_weight_decay:
                if re.search(r, param_name) is not None:
                    return False
        return True

Adam enables L2 weight decay and clip_by_global_norm on gradients. Just adding the square of the weights to the
loss function is *not* the correct way of using L2 regularization/weight decay with Adam, since that will interact
with the m and v parameters in strange ways as shown in [Decoupled Weight Decay
Regularization](https://huggingface.co/papers/1711.05101).

Instead we want to decay the weights in a manner that doesn't interact with the m/v parameters. This is equivalent
to adding the square of the weights to the loss with plain (non-momentum) SGD.

Args:
    learning_rate (`Union[float, LearningRateSchedule]`, *optional*, defaults to 0.001):
        The learning rate to use or a schedule.
    beta_1 (`float`, *optional*, defaults to 0.9):
        The beta1 parameter in Adam, which is the exponential decay rate for the 1st momentum estimates.
    beta_2 (`float`, *optional*, defaults to 0.999):
        The beta2 parameter in Adam, which is the exponential decay rate for the 2nd momentum estimates.
    epsilon (`float`, *optional*, defaults to 1e-07):
        The epsilon parameter in Adam, which is a small constant for numerical stability.
    amsgrad (`bool`, *optional*, defaults to `False`):
        Whether to apply AMSGrad variant of this algorithm or not, see [On the Convergence of Adam and
        Beyond](https://huggingface.co/papers/1904.09237).
    weight_decay_rate (`float`, *optional*, defaults to 0.0):
        The weight decay to apply.
    include_in_weight_decay (`List[str]`, *optional*):
        List of the parameter names (or re patterns) to apply weight decay to. If none is passed, weight decay is
        applied to all parameters by default (unless they are in `exclude_from_weight_decay`).
    exclude_from_weight_decay (`List[str]`, *optional*):
        List of the parameter names (or re patterns) to exclude from applying weight decay to. If a
        `include_in_weight_decay` is passed, the names in it will supersede this list.
    name (`str`, *optional*, defaults to `"AdamWeightDecay"`):
        Optional name for the operations created when applying gradients.
    kwargs (`Dict[str, Any]`, *optional*):
        Keyword arguments. Allowed to be {`clipnorm`, `clipvalue`, `lr`, `decay`}. `clipnorm` is clip gradients by
        norm; `clipvalue` is clip gradients by value, `decay` is included for backward compatibility to allow time
        inverse decay of learning rate. `lr` is included for backward compatibility, recommended to use
        `learning_rate` instead.


def from_text_vision_configs(cls, text_config: AltCLIPTextConfig, vision_config: AltCLIPVisionConfig, **kwargs):
    return cls(text_config=text_config.to_dict(), vision_config=vision_config.to_dict(), **kwargs)

Instantiate a [`AltCLIPConfig`] (or a derived class) from altclip text model configuration and altclip vision
model configuration.

Returns:
    [`AltCLIPConfig`]: An instance of a configuration object


def group_entities(self, entities: List[dict]) -> List[dict]:
    entity_groups = []
    entity_group_disagg = []
    for entity in entities:
        if not entity_group_disagg:
            entity_group_disagg.append(entity)
            continue
        bi, tag = self.get_tag(entity['entity'])
        last_bi, last_tag = self.get_tag(entity_group_disagg[-1]['entity'])
        if tag == last_tag and bi != 'B':
            entity_group_disagg.append(entity)
        else:
            entity_groups.append(self.group_sub_entities(entity_group_disagg))
            entity_group_disagg = [entity]
    if entity_group_disagg:
        entity_groups.append(self.group_sub_entities(entity_group_disagg))
    return entity_groups

Find and group together the adjacent tokens with the same entity predicted.

Args:
    entities (`dict`): The entities predicted by the pipeline.


def eq(left: Any, right: Any) -> bool:
    if left is right:
        return True
    if isinstance(left, list) and isinstance(right, list) or (isinstance(left, tuple) and isinstance(right, tuple)):
        if len(left) != len(right):
            return False
        for x, y in zip(left, right):
            if ne(x, y):
                return False
        return True
    elif isinstance(left, dict):
        if not isinstance(right, dict) or len(left) != len(right) or set(left.keys()) != set(right.keys()):
            return False
        left_items = left.sym_items if isinstance(left, Symbolic) else left.items
        right_item = right.sym_getattr if isinstance(right, Symbolic) else right.__getitem__
        for k, v in left_items():
            if ne(v, right_item(k)):
                return False
        return True
    elif hasattr(left, 'sym_eq') and (not inspect.isclass(left)) and (left.sym_eq.__code__ is not Symbolic.sym_eq.__code__):
        return left.sym_eq(right)
    elif hasattr(right, 'sym_eq') and (not inspect.isclass(right)) and (right.sym_eq.__code__ is not Symbolic.sym_eq.__code__):
        return right.sym_eq(left)
    return pg_typing.callable_eq(left, right)

Compares if two values are equal. Use symbolic equality if possible.

Example::

  @pg.members([
    ('x', pg.typing.Any())
  ])
  class A(pg.Object):
    def sym_eq(self, right):
      if super().sym_eq(right):
        return True
      return pg.eq(self.x, right)

  class B:
    pass

  assert pg.eq(1, 1)
  assert pg.eq(A(1), A(1))
  # This is True since A has override `sym_eq`.
  assert pg.eq(A(1), 1)
  # Objects of B are compared by references.
  assert not pg.eq(A(B()), A(B()))

Args:
  left: The left-hand value to compare.
  right: The right-hand value to compare.

Returns:
  True if left and right is equal or symbolically equal. Otherwise False.


def _set_details(self, content):
    try:
        self.details = str(content)
    except UnicodeEncodeError:
        logging.error('Unable to decode "%s" in Py3, encoding in utf-8.', content)
        self.details = content.encode('utf-8')

Sets the `details` field.

Args:
  content: the content to extract details from.


def write_dirpath(dirpath, strategy):
    if strategy is None:
        strategy = distribute_lib.get_strategy()
    if strategy is None:
        return dirpath
    if not strategy.extended._in_multi_worker_mode():
        return dirpath
    if strategy.extended.should_checkpoint:
        return dirpath
    return _get_temp_dir(dirpath, strategy)

Returns the writing dir that should be used to save file distributedly.

`dirpath` would be created if it doesn't exist.

Args:
  dirpath: Original dirpath that would be used without distribution.
  strategy: The tf.distribute strategy object currently used.

Returns:
  The writing dir path that should be used to save with distribution.


def load_from_saved_model(saved_model_path, custom_objects=None):
    warnings.warn('`tf.keras.experimental.load_from_saved_model` is deprecatedand will be removed in a future version. Please switch to `tf.keras.models.load_model`.')
    model_json_filepath = os.path.join(compat.as_bytes(saved_model_path), compat.as_bytes(constants.ASSETS_DIRECTORY), compat.as_bytes(SAVED_MODEL_FILENAME_JSON))
    with gfile.Open(model_json_filepath, 'r') as f:
        model_json = f.read()
    model = model_config.model_from_json(model_json, custom_objects=custom_objects)
    checkpoint_prefix = os.path.join(compat.as_text(saved_model_path), compat.as_text(constants.VARIABLES_DIRECTORY), compat.as_text(constants.VARIABLES_FILENAME))
    model.load_weights(checkpoint_prefix)
    return model

Loads a keras Model from a SavedModel created by `export_saved_model()`.

This function reinstantiates model state by:
1) loading model topology from json (this will eventually come
   from metagraph).
2) loading model weights from checkpoint.

Example:

```python
import tensorflow as tf

# Create a tf.keras model.
model = tf.keras.Sequential()
model.add(tf.keras.layers.Dense(1, input_shape=[10]))
model.summary()

# Save the tf.keras model in the SavedModel format.
path = '/tmp/simple_keras_model'
tf.keras.experimental.export_saved_model(model, path)

# Load the saved keras model back.
new_model = tf.keras.experimental.load_from_saved_model(path)
new_model.summary()
```

Args:
  saved_model_path: a string specifying the path to an existing SavedModel.
  custom_objects: Optional dictionary mapping names
      (strings) to custom classes or functions to be
      considered during deserialization.

Returns:
  a keras.Model instance.


def _inference_forward(self, z: torch.Tensor, mask: Optional[torch.Tensor]=None, inplace_chunk_size: Optional[int]=None, with_add: bool=True):
    if mask is None:
        mask = z.new_ones(z.shape[:-1])
    mask = mask.unsqueeze(-1)

    def compute_projection_helper(pair, mask, a=True):
        if a:
            linear_g = self.linear_a_g
            linear_p = self.linear_a_p
        else:
            linear_g = self.linear_b_g
            linear_p = self.linear_b_p
        pair = self.layer_norm_in(pair)
        p = linear_g(pair)
        p.sigmoid_()
        p *= linear_p(pair)
        p *= mask
        p = permute_final_dims(p, (2, 0, 1))
        return p

    def compute_projection(pair, mask, a=True, chunked=True):
        need_transpose = self._outgoing ^ a
        if not chunked:
            p = compute_projection_helper(pair, mask, a)
            if need_transpose:
                p = p.transpose(-1, -2)
        else:
            linear_g = self.linear_a_g if a else self.linear_b_g
            c = linear_g.bias.shape[-1]
            out_shape = pair.shape[:-3] + (c,) + pair.shape[-3:-1]
            p = pair.new_zeros(out_shape)
            for i in range(0, pair.shape[-3], inplace_chunk_size):
                pair_chunk = pair[..., i:i + inplace_chunk_size, :, :]
                pair_chunk = compute_projection_helper(pair[..., i:i + inplace_chunk_size, :, :], mask[..., i:i + inplace_chunk_size, :, :], a)
                if need_transpose:
                    pair_chunk = pair_chunk.transpose(-1, -2)
                    p[..., i:i + inplace_chunk_size] = pair_chunk
                else:
                    p[..., i:i + inplace_chunk_size, :] = pair_chunk
                del pair_chunk
        return p
    a = compute_projection(z, mask, True, chunked=True)
    if inplace_chunk_size is not None:
        n = a.shape[-1]
        half_n = n // 2 + n % 2
        row_dim = -3
        col_dim = -2
        b_chunk_dim = row_dim if self._outgoing else col_dim

        def empty_slicer(t):
            return [slice(None) for _ in t.shape]

        def slice_tensor(t, start, end, dim):
            s = empty_slicer(t)
            s[dim] = slice(start, end)
            return t[s]

        def flip_z_cache_(z_cache, z):
            quadrant_3 = slice_tensor(z_cache, half_n, None, row_dim)
            z_cache = z_cache.transpose(row_dim, col_dim)
            z_cache = z_cache[..., :n // 2, :, :]
            first_half_slicer = empty_slicer(z_cache)
            first_half_slicer[col_dim] = slice(0, half_n)
            z_cache[first_half_slicer] = quadrant_3
            quadrant_4 = slice_tensor(z, half_n, None, row_dim)
            quadrant_4 = slice_tensor(quadrant_4, half_n, None, col_dim)
            quadrant_3_slicer = empty_slicer(z_cache)
            quadrant_3_slicer[col_dim] = slice(half_n, None)
            z_cache[quadrant_3_slicer] = quadrant_4
            return z_cache
        z_cache_shape = list(z.shape)
        z_cache_shape[col_dim] = half_n
        z_cache = z.new_zeros(z_cache_shape)
        z_cache_slicer = empty_slicer(z_cache)
        z_cache_slicer[col_dim] = slice(0, half_n)
        z_cache.copy_(z[z_cache_slicer])
        z_cache_rotated = False
        i_range = list(range(0, half_n, inplace_chunk_size))
        initial_offsets = [i_2 - i_1 for i_1, i_2 in zip(i_range, i_range[1:] + [half_n])]
        after_half = list(range(half_n, n, inplace_chunk_size))
        after_half_offsets = [inplace_chunk_size for _ in after_half]
        combined_range_with_offsets = zip(i_range + after_half, initial_offsets + after_half_offsets)
        for i, offset in combined_range_with_offsets:
            if not z_cache_rotated and i >= half_n:
                z_cache = flip_z_cache_(z_cache, z)
                z_cache_rotated = True
            z_chunk_b = slice_tensor(z, i, i + offset, b_chunk_dim)
            mask_chunk = slice_tensor(mask, i, i + offset, b_chunk_dim)
            z_chunk_b = z_chunk_b.clone()
            if b_chunk_dim == col_dim:
                z_chunk_b = slice_tensor(z, i, i + offset, col_dim)
            elif not z_cache_rotated:
                z_chunk_slicer = empty_slicer(z_chunk_b)
                z_chunk_slicer[col_dim] = slice(0, half_n)
                z_chunk_b[z_chunk_slicer] = slice_tensor(z_cache, i, i + offset, row_dim)
            else:
                z_cache_offset = i - half_n
                z_chunk_b = slice_tensor(z_cache, z_cache_offset, z_cache_offset + offset, row_dim)
            b_chunk = compute_projection(z_chunk_b, mask_chunk, a=False, chunked=False)
            del z_chunk_b
            x_chunk = torch.matmul(a, b_chunk)
            x_chunk = permute_final_dims(x_chunk, (1, 2, 0))
            x_chunk = self.layer_norm_out(x_chunk)
            x_chunk = self.linear_z(x_chunk)
            z_chunk_g = slice_tensor(z, i, i + offset, col_dim)
            g_chunk = self.linear_g(self.layer_norm_in(z_chunk_g))
            g_chunk.sigmoid_()
            del z_chunk_g
            x_chunk *= g_chunk
            z_slicer = empty_slicer(z)
            z_slicer[col_dim] = slice(i, i + offset)
            if with_add:
                z[z_slicer] += x_chunk
            else:
                z[z_slicer] = x_chunk
    else:
        b = compute_projection(z, mask, False, False)
        x = torch.matmul(a, b)
        x = self.layer_norm_out(x)
        x = self.linear_z(x)
        g = self.linear_g(z)
        g.sigmoid_()
        x *= g
        if with_add:
            z += x
        else:
            z = x
    return z

Args:
    z:
        A [*, N, N, C_z] pair representation
    mask:
        A [*, N, N] pair mask
    inplace_chunk_size:
        Size of chunks used in the main computation. Increase to trade memory for speed.
    with_add:
        If True, z is overwritten with (z + update). Otherwise, it is overwritten with (update).
Returns:
    A reference to the overwritten z

More memory-efficient, inference-only version of the forward function. Uses in-place operations, fusion of the
addition that happens after this module in the Evoformer, a smidge of recomputation, and a cache of overwritten
values to lower peak memory consumption of this module from 5x the size of the input tensor z to 2.5x its size.
Useful for inference on extremely long sequences.

It works as follows. We will make reference to variables used in the default forward implementation below.
Naively, triangle multiplication attention requires the manifestation of 5 tensors the size of z: 1) z, the
"square" input tensor, 2) a, the first projection of z, 3) b, the second projection of b, 4) g, a z-sized mask,
and 5) a z-sized tensor for intermediate computations. For large N, this is prohibitively expensive; for
N=4000, for example, z is more than 8GB alone. To avoid this problem, we compute b, g, and all intermediate
tensors in small chunks, noting that the chunks required to compute a chunk of the output depend only on the
tensor a and corresponding vertical and horizontal chunks of z. This suggests an algorithm that loops over
pairs of chunks of z: hereafter "columns" and "rows" of z, even though each "column" and "row" in fact contains
inplace_chunk_size contiguous true columns and rows of z. Writing output chunks to a new tensor would bring
total memory consumption down to 3x the size of z. However, more memory can be saved by writing output chunks
directly to z in-place. WLOG, we choose to write output chunks vertically, overwriting the ith "column" of z at
the end of the ith iteration of the main loop. Despite this overwriting, the ith column is always one column
ahead of previously overwritten columns and can be recovered directly from z. After the first iteration,
however, the ith row of z is always at least partially overwritten. For this reason, we introduce the z-cache,
a tensor one-half the size of z. The z-cache initially contains the left half (2nd and 3rd quadrants) of z. For
0 < i < N/2, the missing left part of the ith row of z is recovered from this cache at the beginning of the ith
iteration. Once i exceeds n/2, the cache is "reoriented" to encompass the 3rd and 4th quadrants of z instead.
Though the 3rd quadrant of the original z is entirely overwritten at this point, it can be recovered from the
z-cache itself. Thereafter, the ith row of z can be recovered in its entirety from the reoriented z-cache.
After the final iteration, z has been completely overwritten and contains the triangular multiplicative update.
If with_add is True, it instead contains the sum of z and the triangular multiplicative update. In either case,
peak memory consumption is just 2.5x the size of z, disregarding memory used for chunks and other small
variables.


def get_tf_dataset(self):
    raise NotImplementedError

Get a `tf.data.Dataset` instance for the DataAdapter.

Note that the dataset returned does not repeat for epoch, so caller
might need to create new iterator for the same dataset at the beginning
of the epoch. This behavior might change in the future.

Returns:
    A `tf.data.Dataset`. Caller might use the dataset in different
    context, e.g. iter(dataset) in eager to get the value directly, or
    in graph mode, provide the iterator tensor to Keras model function.


class Node:

    def __init__(self, operation, call_args=None, call_kwargs=None, outputs=None):
        self.operation = operation
        self.arguments = SymbolicArguments(*call_args, **call_kwargs)
        self.outputs = [] if outputs is None else tree.flatten(outputs)
        for x in self.outputs:
            if not isinstance(x, KerasTensor):
                raise ValueError(f'All operation outputs must be tensors. Operation {operation} returned a non-tensor. Non-tensor received: {x}')
        zero_history = any((not x.record_history for x in self.arguments.keras_tensors))
        if not zero_history:
            for tensor in self.arguments.keras_tensors:
                if not hasattr(tensor, '_keras_history'):
                    tensor._keras_history = KerasHistory(operation=None, node_index=0, tensor_index=0)
        self.operation._inbound_nodes.append(self)
        for kt in self.arguments.keras_tensors:
            inbound_op = kt._keras_history.operation
            if inbound_op is not None:
                inbound_op._outbound_nodes.append(self)
        if not zero_history:
            node_index = len(self.operation._inbound_nodes) - 1
            for i, tensor in enumerate(self.outputs):
                tensor._keras_history = KerasHistory(operation=operation, node_index=node_index, tensor_index=i)
        self.is_input = not self.arguments.keras_tensors

    def __repr__(self):
        return f'<Node operation={self.operation.name}, id={id(self)}>'

    @property
    def input_tensors(self):
        return self.arguments.keras_tensors

    @property
    def output_tensors(self):
        return self.outputs

    @property
    def parent_nodes(self):
        """The parent `Node`s.

        Returns:
            all the `Node`s whose output this node immediately depends on.
        """
        node_deps = []
        for kt in self.arguments.keras_tensors:
            op = kt._keras_history.operation
            node_index = kt._keras_history.node_index
            if op is not None:
                node_deps.append(op._inbound_nodes[node_index])
        return node_deps

A `Node` describes an operation `__call__()` event.

A Keras Function is a DAG with `Node` instances as nodes, and
`KerasTensor` instances as edges. Nodes aren't `Operation` instances,
because a single operation could be called multiple times, which would
result in graph cycles.

A `__call__()` event involves input tensors (and other input arguments),
the operation that was called, and the resulting output tensors.
A `Node` will include all this information.

Since a single `Operation` could be called multiple times,
the `Node` instances are stored on operations as a list.
Each time an operation is called, a node is added to `op._inbound_nodes`.
Each time the output of an operation is used by another operation,
a node is added to `op._outbound_nodes`.

Every `KerasTensor` instance has a `KerasHistory` object attached,
which tracks the `Node` that records the `__call__()` event that created
the tensor. By recursively walking through `Node` instances
via the `KerasHistory` metadata of `KerasTensor` instances, once can
retrieve the entire DAG of a Keras Function.

Args:
    operation: The Operation that was called in the `op.__call__()`
        event that this node represents.
    call_args: The positional arguments the operation was called with.
    call_kwargs: The keyword arguments the operation was called with.
    outputs: The output tensors of the `op.__call__()` call.