def clone_model(model, input_tensors=None, clone_function=None):
    with generic_utils.DisableSharedObjectScope():
        if clone_function is None:
            clone_function = _clone_layer
        if isinstance(model, Sequential):
            return _clone_sequential_model(model, input_tensors=input_tensors, layer_fn=clone_function)
        else:
            return _clone_functional_model(model, input_tensors=input_tensors, layer_fn=clone_function)

Clone a Functional or Sequential `Model` instance.

Model cloning is similar to calling a model on new inputs,
except that it creates new layers (and thus new weights) instead
of sharing the weights of the existing layers.

Note that
`clone_model` will not preserve the uniqueness of shared objects within the
model (e.g. a single variable attached to two distinct layers will be
restored as two separate variables).

Args:
    model: Instance of `Model`
        (could be a Functional model or a Sequential model).
    input_tensors: optional list of input tensors or InputLayer objects
        to build the model upon. If not provided,
        new `Input` objects will be created.
    clone_function: Callable to be used to clone each layer in the target
        model (except `InputLayer` instances). It takes as argument the layer
        instance to be cloned, and returns the corresponding layer instance to
        be used in the model copy. If unspecified, this callable defaults to
        the following serialization/deserialization function:
        `lambda layer: layer.__class__.from_config(layer.get_config())`.
        By passing a custom callable, you can customize your copy of the
        model, e.g. by wrapping certain layers of interest (you might want to
        replace all `LSTM` instances with equivalent
        `Bidirectional(LSTM(...))` instances, for example).

Returns:
  An instance of `Model` reproducing the behavior
  of the original model, on top of new inputs tensors,
  using newly instantiated weights. The cloned model may behave
  differently from the original model if a custom `clone_function`
  modifies the layer.

Example:

```python
# Create a test Sequential model.
model = keras.Sequential([
    keras.Input(shape=(728,)),
    keras.layers.Dense(32, activation='relu'),
    keras.layers.Dense(1, activation='sigmoid'),
])
# Create a copy of the test model (with freshly initialized weights).
new_model = clone_model(model)
```

Note that subclassed models cannot be cloned, since their internal
layer structure is not known. To achieve equivalent functionality
as `clone_model` in the case of a subclassed model, simply make sure
that the model class implements `get_config()`
(and optionally `from_config()`), and call:

```python
new_model = model.__class__.from_config(model.get_config())
```


class GPTBigCodeConfig(PretrainedConfig):
    model_type = 'gpt_bigcode'
    keys_to_ignore_at_inference = ['past_key_values']
    attribute_map = {'hidden_size': 'n_embd', 'max_position_embeddings': 'n_positions', 'num_attention_heads': 'n_head', 'num_hidden_layers': 'n_layer'}

    def __init__(self, vocab_size=50257, n_positions=1024, n_embd=768, n_layer=12, n_head=12, n_inner=None, activation_function='gelu_pytorch_tanh', resid_pdrop=0.1, embd_pdrop=0.1, attn_pdrop=0.1, layer_norm_epsilon=1e-05, initializer_range=0.02, scale_attn_weights=True, use_cache=True, bos_token_id=50256, eos_token_id=50256, attention_softmax_in_fp32=True, scale_attention_softmax_in_fp32=True, multi_query=True, **kwargs):
        self.vocab_size = vocab_size
        self.n_positions = n_positions
        self.n_embd = n_embd
        self.n_layer = n_layer
        self.n_head = n_head
        self.n_inner = n_inner
        self.activation_function = activation_function
        self.resid_pdrop = resid_pdrop
        self.embd_pdrop = embd_pdrop
        self.attn_pdrop = attn_pdrop
        self.layer_norm_epsilon = layer_norm_epsilon
        self.initializer_range = initializer_range
        self.scale_attn_weights = scale_attn_weights
        self.use_cache = use_cache
        self.attention_softmax_in_fp32 = attention_softmax_in_fp32
        self.scale_attention_softmax_in_fp32 = scale_attention_softmax_in_fp32
        self.multi_query = multi_query
        self.bos_token_id = bos_token_id
        self.eos_token_id = eos_token_id
        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)

This is the configuration class to store the configuration of a [`GPTBigCodeModel`]. It is used to instantiate a
GPTBigCode model according to the specified arguments, defining the model architecture. Instantiating a
configuration with the defaults will yield a similar configuration to that of the GPTBigCode
[gpt_bigcode](https://huggingface.co/gpt_bigcode) architecture.

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

Args:
    vocab_size (`int`, *optional*, defaults to 50257):
        Vocabulary size of the GPT-2 model. Defines the number of different tokens that can be represented by the
        `inputs_ids` passed when calling [`GPTBigCodeModel`].
    n_positions (`int`, *optional*, defaults to 1024):
        The maximum sequence length that this model might ever be used with. Typically set this to something large
        just in case (e.g., 512 or 1024 or 2048).
    n_embd (`int`, *optional*, defaults to 768):
        Dimensionality of the embeddings and hidden states.
    n_layer (`int`, *optional*, defaults to 12):
        Number of hidden layers in the Transformer encoder.
    n_head (`int`, *optional*, defaults to 12):
        Number of attention heads for each attention layer in the Transformer encoder.
    n_inner (`int`, *optional*, defaults to None):
        Dimensionality of the inner feed-forward layers. `None` will set it to 4 times n_embd
    activation_function (`str`, *optional*, defaults to `"gelu_pytorch_tanh"`):
        Activation function, to be selected in the list `["relu", "silu", "gelu", "tanh", "gelu_new",
        "gelu_pytorch_tanh"]`.
    resid_pdrop (`float`, *optional*, defaults to 0.1):
        The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
    embd_pdrop (`float`, *optional*, defaults to 0.1):
        The dropout ratio for the embeddings.
    attn_pdrop (`float`, *optional*, defaults to 0.1):
        The dropout ratio for the attention.
    layer_norm_epsilon (`float`, *optional*, defaults to 1e-5):
        The epsilon to use in the layer normalization layers.
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    scale_attn_weights (`bool`, *optional*, defaults to `True`):
        Scale attention weights by dividing by sqrt(hidden_size)..
    use_cache (`bool`, *optional*, defaults to `True`):
        Whether or not the model should return the last key/values attentions (not used by all models).
    attention_softmax_in_fp32 (`bool`, *optional*, defaults to `True`):
        Whether to call the fused softmax in float32.
    scale_attention_softmax_in_fp32 (`bool`, *optional*, defaults to `True`):
        Whether to scale the attention softmax in float32.
    attention_type (`bool`, *optional*, defaults to `True`):
        Whether to use Multi-Query Attion (`True`) or Multi-Head Attention (`False`).
Example:

```python
>>> from transformers import GPTBigCodeConfig, GPTBigCodeModel

>>> # Initializing a GPTBigCode configuration
>>> configuration = GPTBigCodeConfig()

>>> # Initializing a model (with random weights) from the configuration
>>> model = GPTBigCodeModel(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config
```


def get_np_doc_form():
    return _np_doc_form

Gets the form of the original numpy docstrings.

Returns:
  See `set_np_doc_form` for the list of valid values.


def create_network_inputs(self, past_values: torch.Tensor, past_time_features: torch.Tensor, static_categorical_features: Optional[torch.Tensor]=None, static_real_features: Optional[torch.Tensor]=None, past_observed_mask: Optional[torch.Tensor]=None, future_values: Optional[torch.Tensor]=None, future_time_features: Optional[torch.Tensor]=None) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    time_feat = torch.cat((past_time_features[:, self._past_length - self.config.context_length:, ...], future_time_features), dim=1) if future_values is not None else past_time_features[:, self._past_length - self.config.context_length:, ...]
    if past_observed_mask is None:
        past_observed_mask = torch.ones_like(past_values)
    context = past_values[:, -self.config.context_length:]
    observed_context = past_observed_mask[:, -self.config.context_length:]
    _, loc, scale = self.scaler(context, observed_context)
    inputs = (torch.cat((past_values, future_values), dim=1) - loc) / scale if future_values is not None else (past_values - loc) / scale
    log_abs_loc = loc.abs().log1p() if self.config.input_size == 1 else loc.squeeze(1).abs().log1p()
    log_scale = scale.log() if self.config.input_size == 1 else scale.squeeze(1).log()
    static_feat = torch.cat((log_abs_loc, log_scale), dim=1)
    if static_real_features is not None:
        static_feat = torch.cat((static_real_features, static_feat), dim=1)
    if static_categorical_features is not None:
        embedded_cat = self.embedder(static_categorical_features)
        static_feat = torch.cat((embedded_cat, static_feat), dim=1)
    expanded_static_feat = static_feat.unsqueeze(1).expand(-1, time_feat.shape[1], -1)
    features = torch.cat((expanded_static_feat, time_feat), dim=-1)
    subsequences_length = self.config.context_length + self.config.prediction_length if future_values is not None else self.config.context_length
    lagged_sequence = self.get_lagged_subsequences(sequence=inputs, subsequences_length=subsequences_length)
    lags_shape = lagged_sequence.shape
    reshaped_lagged_sequence = lagged_sequence.reshape(lags_shape[0], lags_shape[1], -1)
    if reshaped_lagged_sequence.shape[1] != time_feat.shape[1]:
        raise ValueError(f'input length {reshaped_lagged_sequence.shape[1]} and time feature lengths {time_feat.shape[1]} does not match')
    return (reshaped_lagged_sequence, features, loc, scale, static_feat)

Creates the inputs for the network given the past and future values, time features, and static features.

Args:
    past_values (`torch.Tensor`):
        A tensor of shape `(batch_size, past_length, input_size)` containing the past values.
    past_time_features (`torch.Tensor`):
        A tensor of shape `(batch_size, past_length, num_features)` containing the past time features.
    static_categorical_features (`Optional[torch.Tensor]`):
        An optional tensor of shape `(batch_size, num_categorical_features)` containing the static categorical
        features.
    static_real_features (`Optional[torch.Tensor]`):
        An optional tensor of shape `(batch_size, num_real_features)` containing the static real features.
    past_observed_mask (`Optional[torch.Tensor]`):
        An optional tensor of shape `(batch_size, past_length, input_size)` containing the mask of observed
        values in the past.
    future_values (`Optional[torch.Tensor]`):
        An optional tensor of shape `(batch_size, future_length, input_size)` containing the future values.

Returns:
    A tuple containing the following tensors:
    - reshaped_lagged_sequence (`torch.Tensor`): A tensor of shape `(batch_size, sequence_length, num_lags *
      input_size)` containing the lagged subsequences of the inputs.
    - features (`torch.Tensor`): A tensor of shape `(batch_size, sequence_length, num_features)` containing the
      concatenated static and time features.
    - loc (`torch.Tensor`): A tensor of shape `(batch_size, input_size)` containing the mean of the input
      values.
    - scale (`torch.Tensor`): A tensor of shape `(batch_size, input_size)` containing the std of the input
      values.
    - static_feat (`torch.Tensor`): A tensor of shape `(batch_size, num_static_features)` containing the
      concatenated static features.


def _CompletionsFromArgs(fn_args):
    completions = []
    for arg in fn_args:
        arg = arg.replace('_', '-')
        completions.append(f'--{arg}')
    return completions

Takes a list of fn args and returns a list of the fn's completion strings.

Args:
  fn_args: A list of the args accepted by a function.
Returns:
  A list of possible completion strings for that function.


def forward(self, pixel_values: torch.FloatTensor, prompt_depth: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple[torch.Tensor], DepthEstimatorOutput]:
    loss = None
    if labels is not None:
        raise NotImplementedError('Training is not implemented yet')
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    outputs = self.backbone.forward_with_filtered_kwargs(pixel_values, output_hidden_states=output_hidden_states, output_attentions=output_attentions)
    hidden_states = outputs.feature_maps
    _, _, height, width = pixel_values.shape
    patch_size = self.config.patch_size
    patch_height = height // patch_size
    patch_width = width // patch_size
    if prompt_depth is not None:
        batch_size = prompt_depth.shape[0]
        depth_min = torch.min(prompt_depth.reshape(batch_size, -1), dim=1).values
        depth_max = torch.max(prompt_depth.reshape(batch_size, -1), dim=1).values
        depth_min, depth_max = (depth_min.view(batch_size, 1, 1, 1), depth_max.view(batch_size, 1, 1, 1))
        prompt_depth = (prompt_depth - depth_min) / (depth_max - depth_min)
    hidden_states = self.neck(hidden_states, patch_height, patch_width, prompt_depth=prompt_depth)
    predicted_depth = self.head(hidden_states, patch_height, patch_width)
    if prompt_depth is not None:
        depth_min = depth_min.squeeze(1).to(predicted_depth.device)
        depth_max = depth_max.squeeze(1).to(predicted_depth.device)
        predicted_depth = predicted_depth * (depth_max - depth_min) + depth_min
    if not return_dict:
        if output_hidden_states:
            output = (predicted_depth,) + outputs[1:]
        else:
            output = (predicted_depth,) + outputs[2:]
        return (loss,) + output if loss is not None else output
    return DepthEstimatorOutput(loss=loss, predicted_depth=predicted_depth, hidden_states=outputs.hidden_states if output_hidden_states else None, attentions=outputs.attentions)

prompt_depth (`torch.FloatTensor` of shape `(batch_size, 1, height, width)`, *optional*):
    Prompt depth is the sparse or low-resolution depth obtained from multi-view geometry or a
    low-resolution depth sensor. It generally has shape (height, width), where height
    and width can be smaller than those of the images. It is optional and can be None, which means no prompt depth
    will be used. If it is None, the output will be a monocular relative depth.
    The values are recommended to be in meters, but this is not necessary.

Example:

```python
>>> from transformers import AutoImageProcessor, AutoModelForDepthEstimation
>>> import torch
>>> import numpy as np
>>> from PIL import Image
>>> import requests

>>> url = "https://github.com/DepthAnything/PromptDA/blob/main/assets/example_images/image.jpg?raw=true"
>>> image = Image.open(requests.get(url, stream=True).raw)

>>> image_processor = AutoImageProcessor.from_pretrained("depth-anything/prompt-depth-anything-vits-hf")
>>> model = AutoModelForDepthEstimation.from_pretrained("depth-anything/prompt-depth-anything-vits-hf")

>>> prompt_depth_url = "https://github.com/DepthAnything/PromptDA/blob/main/assets/example_images/arkit_depth.png?raw=true"
>>> prompt_depth = Image.open(requests.get(prompt_depth_url, stream=True).raw)

>>> # prepare image for the model
>>> inputs = image_processor(images=image, return_tensors="pt", prompt_depth=prompt_depth)

>>> with torch.no_grad():
...     outputs = model(**inputs)

>>> # interpolate to original size
>>> post_processed_output = image_processor.post_process_depth_estimation(
...     outputs,
...     target_sizes=[(image.height, image.width)],
... )

>>> # visualize the prediction
>>> predicted_depth = post_processed_output[0]["predicted_depth"]
>>> depth = predicted_depth * 1000.
>>> depth = depth.detach().cpu().numpy()
>>> depth = Image.fromarray(depth.astype("uint16")) # mm
```


def get_dict(self):
    self.is_valid()
    return self._get_dict()

Returns the internal-API dictionary representing the
:class:`DisplayDataItem`.

Returns:
  Dict[str, Any]: A dictionary. The internal-API dictionary representing
  the :class:`DisplayDataItem`.

Raises:
  ValueError: if the item is not valid.


def convert_date_to_duration(date: Timestamp) -> NormalizedDuration:
    if isinstance(date, float):
        return date
    if isinstance(date, int):
        return float(date)
    if isinstance(date, np.datetime64):
        return convert_numpy_datetime64_to_duration(date)
    if isinstance(date, datetime.datetime):
        return convert_datetime_to_duration(date)
    if isinstance(date, datetime.date):
        return convert_datetime_date_to_duration(date)
    raise TypeError(f'Unsupported type: {type(date)}')

Converts date value to a number representing the Unix timestamp.

If a float or int, it is returned as float.
If a date, it is converted to a Unix timestamp (number of seconds from Unix
epoch).

Args:
    date: Date to convert.

Returns:
    Unix timestamp (seconds elapsed from unix epoch).

Raises:
    TypeError: unsupported type. Supported types are:
        - np.datetime64
        - datetime.datetime


def _process_stack_frames(self):
    stack_frames = tf_stack.extract_stack()
    stack_frame_ids = []
    writer = None
    for file_path, lineno, func, _ in stack_frames:
        abs_path = os.path.abspath(file_path)
        if (abs_path, lineno, func) in self._stack_frame_to_id:
            stack_frame_ids.append(self._stack_frame_to_id[abs_path, lineno, func])
            continue
        with self._stack_frame_to_id_lock:
            if (abs_path, lineno, func) not in self._stack_frame_to_id:
                stack_frame_id = _get_id()
                self._stack_frame_to_id[abs_path, lineno, func] = stack_frame_id
                file_index = self._write_source_file_content(abs_path)
                file_line_col = graph_debug_info_pb2.GraphDebugInfo.FileLineCol(file_index=file_index, line=lineno, func=func)
                stack_frame_with_id = debug_event_pb2.StackFrameWithId(id=stack_frame_id, file_line_col=file_line_col)
                writer = self.get_writer()
                writer.WriteStackFrameWithId(stack_frame_with_id)
            stack_frame_ids.append(self._stack_frame_to_id[abs_path, lineno, func])
    code_location = debug_event_pb2.CodeLocation(host_name=self._hostname, stack_frame_ids=stack_frame_ids)
    return code_location

Process stack frames.

Send the content of source-files, on a best-effort basis.

Returns:
  A list of stack frame IDs.


def create_position_ids_from_inputs_embeds(self, inputs_embeds):
    input_shape = inputs_embeds.size()[:-1]
    sequence_length = input_shape[1]
    position_ids = torch.arange(self.padding_idx + 1, sequence_length + self.padding_idx + 1, dtype=torch.long, device=inputs_embeds.device)
    return position_ids.unsqueeze(0).expand(input_shape)

We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.

Args:
    inputs_embeds: torch.Tensor

Returns: torch.Tensor


def infer_element_type(elements):
    element_type = typehints.Union[[trivial_inference.instance_to_type(e) for e in elements]]
    return element_type

For internal use only; no backwards-compatibility guarantees.

Infer a Beam type for a list of elements.

Args:
  elements (List[Any]): A list of elements for which the type should be
      inferred.

Returns:
  A Beam type encompassing all elements.


class PromptDepthAnythingFeatureFusionLayer(nn.Module):

    def __init__(self, config: PromptDepthAnythingConfig):
        super().__init__()
        self.projection = nn.Conv2d(config.fusion_hidden_size, config.fusion_hidden_size, kernel_size=1, bias=True)
        self.residual_layer1 = PromptDepthAnythingPreActResidualLayer(config)
        self.residual_layer2 = PromptDepthAnythingPreActResidualLayer(config)
        self.prompt_depth_layer = PromptDepthAnythingLayer(config)

    def forward(self, hidden_state, residual=None, size=None, prompt_depth=None):
        if residual is not None:
            if hidden_state.shape != residual.shape:
                residual = nn.functional.interpolate(residual, size=hidden_state.shape[2:], mode='bilinear', align_corners=False)
            hidden_state = hidden_state + self.residual_layer1(residual)
        hidden_state = self.residual_layer2(hidden_state)
        if prompt_depth is not None:
            prompt_depth = nn.functional.interpolate(prompt_depth, size=hidden_state.shape[2:], mode='bilinear', align_corners=False)
            res = self.prompt_depth_layer(prompt_depth)
            hidden_state = hidden_state + res
        modifier = {'scale_factor': 2} if size is None else {'size': size}
        hidden_state = nn.functional.interpolate(hidden_state, **modifier, mode='bilinear', align_corners=True)
        hidden_state = self.projection(hidden_state)
        return hidden_state

Feature fusion layer, merges feature maps from different stages.

Args:
    config (`[PromptDepthAnythingConfig]`):
        Model configuration class defining the model architecture.


def from_backbone_config(cls, backbone_config: PretrainedConfig, **kwargs):
    return cls(backbone_config=backbone_config, **kwargs)

Instantiate a [`Mask2FormerConfig`] (or a derived class) from a pre-trained backbone model configuration.

Args:
    backbone_config ([`PretrainedConfig`]):
        The backbone configuration.

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


def recipe_bigquery_run_query(config, auth_write, query, legacy):
    bigquery(config, {'auth': auth_write, 'run': {'query': query, 'legacy': legacy}})

Run query on a project.

Args:
  auth_write (authentication) - Credentials used for writing data.
  query (text) - SQL with newlines and all.
  legacy (boolean) - Query type must match table and query format.


def forEach(self) -> 'ColumnExpressionBuilder':
    return ColumnExpressionBuilder(self._builder, self._column_name, self._children, True, True)

The forEach() function.

Unnests the repeated values from a FHIR path. If the FHIR path does not
return a collection, we treat that as a collection with a single value.
Once this function is called, the FHIR path is sealed to be immutable.

Returns:
  A new ColumnExpressionBuilder with needs_unnest set to True.


def _write_cache(step, event_file_suffix=None, **kwargs):
    file_suffix = _TT_EVENT_FILE_SUFFIX
    if event_file_suffix is not None:
        file_suffix = string_ops.string_join([file_suffix, event_file_suffix], separator='.')
    summary_write_ops = []
    summary_writer = summary.create_file_writer_v2(self._parameters.trace_dir, filename_suffix=file_suffix, max_queue=_TT_SUMMARY_MAX_QUEUE)
    graph.add_to_collection(TENSOR_TRACER_SUMMARY_COLLECTION, summary_writer)
    step_value = step[0]
    dt = step_value.dtype
    if dt.__ne__(dtypes.int64) and dt.__ne__(dtypes.uint64) and dt.__ne__(dtypes.float64):
        step_value = math_ops.cast(step_value, dtypes.int64)
    with summary_writer.as_default():
        summary_metadata = summary_pb2.SummaryMetadata(plugin_data=summary_pb2.SummaryMetadata.PluginData(plugin_name=_TT_TENSORBOARD_PLUGIN_NAME))
        for key, value in kwargs.items():
            if not self._parameters.collect_summary_per_core:
                if key == _TT_SUMMARY_TAG and value.shape.as_list()[0] != 1:
                    value = self.aggregate_global_cache(value)
            with ops.control_dependencies([summary_writer.init()]):
                summary_write_ops.append(summary.write(_TT_SUMMARY_TAG + '/' + key + '.' + graph_summary_tag, value, metadata=summary_metadata, step=step_value))
    return control_flow_ops.group(summary_write_ops)

Writes the given caches as tensor summary.

Args:
  step: Step tensor with dimension [num_cores].
  event_file_suffix: Event filename suffix tensor.
  **kwargs: The dictionary of tensors that needs to be written as
    summaries. Key and value pairs within kwargs correspond to the tag
    name, and tensor content that will be written using summary.write.
    The trace_modes that use this function are:
      - summary: In summary mode, kwargs includes a single (tag, content)
      pair which are, _TT_SUMMARY_TAG and a tf.float32 signature_cache
      variable. The dimension of the signature_cache is:
        num_cores x num_traced_tensors x num_signatures.
      - full_tensor_summary: kwargs will include all traced tensors. Tag
      and content correspond to the name of the tensor, and its actual
      content.
Returns:
  A tf.Operation that needs to be executed for the host call dependencies.


def _jvp_helper_wrapper(op_name, attr_tuple, inputs, outputs, tangents, use_batch):
    if use_batch:
        for primal, tangent in zip(inputs, tangents):
            if not tangent.shape.is_compatible_with([None] + primal.shape):
                raise ValueError('Tangent {} was expected to be of shape {} but is instead of shape {}'.format(tangent, [None] + primal.shape, tangent.shape))
        return control_flow_ops.vectorized_map(functools.partial(_jvp_helper, op_name, attr_tuple, inputs, outputs), tangents)
    return _jvp_helper(op_name, attr_tuple, inputs, outputs, tangents)

Computes a batch of Jacobian-vector product for an op.

Args:
  op_name: A string, the type of operation being executed.
  attr_tuple: Attributes of the operation.
  inputs: A flat list of input Tensors to the operation.
  outputs: A flat list of output Tensors from the operation.
  tangents: A flat list of Tensors, compatible with shape `[None] +
    input_shape`.
  use_batch: A bool, True to vetorize over batch of tangents of shape `[None]
    + input_shape`.

Returns:
  A flat list of tangents compatible with `outputs`
  or `[None] + output_shape`.

Raises:
  ValueError: if tangent shapes are not compatible with input shapes.


def assert_shapes_v2(shapes, data=None, summarize=None, message=None, name=None):
    assert_shapes(shapes, data=data, summarize=summarize, message=message, name=name)

Assert tensor shapes and dimension size relationships between tensors.

This Op checks that a collection of tensors shape relationships
satisfies given constraints.

Example:

>>> n = 10
>>> q = 3
>>> d = 7
>>> x = tf.zeros([n,q])
>>> y = tf.ones([n,d])
>>> param = tf.Variable([1.0, 2.0, 3.0])
>>> scalar = 1.0
>>> tf.debugging.assert_shapes([
...  (x, ('N', 'Q')),
...  (y, ('N', 'D')),
...  (param, ('Q',)),
...  (scalar, ()),
... ])

>>> tf.debugging.assert_shapes([
...   (x, ('N', 'D')),
...   (y, ('N', 'D'))
... ])
Traceback (most recent call last):
...
ValueError: ...

If `x`, `y`, `param` or `scalar` does not have a shape that satisfies
all specified constraints, `message`, as well as the first `summarize` entries
of the first encountered violating tensor are printed, and
`InvalidArgumentError` is raised.

Size entries in the specified shapes are checked against other entries by
their __hash__, except:
  - a size entry is interpreted as an explicit size if it can be parsed as an
    integer primitive.
  - a size entry is interpreted as *any* size if it is None or '.'.

If the first entry of a shape is `...` (type `Ellipsis`) or '*' that indicates
a variable number of outer dimensions of unspecified size, i.e. the constraint
applies to the inner-most dimensions only.

Scalar tensors and specified shapes of length zero (excluding the 'inner-most'
prefix) are both treated as having a single dimension of size one.

Args:
  shapes: dictionary with (`Tensor` to shape) items, or a list of
    (`Tensor`, shape) tuples. A shape must be an iterable.
  data: The tensors to print out if the condition is False.  Defaults to error
    message and first few entries of the violating tensor.
  summarize: Print this many entries of the tensor.
  message: A string to prefix to the default message.
  name: A name for this operation (optional).  Defaults to "assert_shapes".

Raises:
  ValueError:  If static checks determine any shape constraint is violated.


def get_gradients(self, loss, params):
    grads = backend.gradients(loss, params)
    if any((g is None for g in grads)):
        raise ValueError('An operation has `None` for gradient. Please make sure that all of your ops have a gradient defined (i.e. are differentiable). Common ops without gradient: backend.argmax, backend.round, backend.eval.')
    if hasattr(self, 'clipnorm'):
        grads = [clip_ops.clip_by_norm(g, self.clipnorm) for g in grads]
    if hasattr(self, 'clipvalue'):
        grads = [clip_ops.clip_by_value(g, -self.clipvalue, self.clipvalue) for g in grads]
    return grads

Returns gradients of `loss` with respect to `params`.

Args:
    loss: Loss tensor.
    params: List of variables.

Returns:
    List of gradient tensors.

Raises:
    ValueError: In case any gradient cannot be computed (e.g. if gradient
      function not implemented).


def fd_solver_forward(self, start_time, end_time, coord_grid, values_grid, one_step_fn=None, boundary_conditions=None, start_step_count=0, num_steps=None, time_step=None, values_transform_fn=None, dtype=None, name=None, **kwargs):
    pde_solver_fn = kwargs.get('pde_solver_fn', fd_solvers.solve_forward)
    backward_second_order, backward_first_order, backward_zeroth_order = _backward_pde_coeffs(self._drift_fn, self._volatility_fn, discounting=None)
    inner_second_order_coeff_fn = lambda t, x: -backward_second_order(t, x)
    inner_first_order_coeff_fn = backward_first_order
    zeroth_order_coeff_fn = backward_zeroth_order
    return pde_solver_fn(start_time=start_time, end_time=end_time, coord_grid=coord_grid, values_grid=values_grid, num_steps=num_steps, start_step_count=start_step_count, time_step=time_step, one_step_fn=one_step_fn, boundary_conditions=boundary_conditions, values_transform_fn=values_transform_fn, inner_second_order_coeff_fn=inner_second_order_coeff_fn, inner_first_order_coeff_fn=inner_first_order_coeff_fn, zeroth_order_coeff_fn=zeroth_order_coeff_fn, dtype=dtype, name=name)

Returns a solver for the Fokker Plank equation of this process.

The Fokker Plank equation (also known as the Kolmogorov Forward equation)
associated to this Ito process is given by:

```None
  dV/dt + d(mean_i(t, x) V) / dx
    - (1/2) d^2(volatility^2(t, x) V) / dx^2 =  = 0
```

with the initial value condition $$V(0, x) = u(x)$$.

This method evolves a spatially discretized solution of the above PDE from
time `t0` to time `t1 < t0` (i.e. backwards in time).
The solution `V(t,x)` is assumed to be discretized on a grid.

This method allows batching of solutions. In this context, batching means
the ability to represent and evolve multiple independent functions `V`
(e.g. V1, V2 ...) simultaneously corresponding to `mean_1, mean_2 ...` and
`volatility_1, volatility_2 ....`.

The evolution of the solution from `t0` to `t1` is often done by
discretizing the differential equation to a difference equation along
the spatial and temporal axes. The temporal discretization is given by a
(sequence of) time steps [dt_1, dt_2, ... dt_k] such that the sum of the
time steps is equal to the total time step `t0 - t1`. If a uniform time
step is used, it may equivalently be specified by stating the number of
steps (n_steps) to take. This method provides both options via the
`time_step` and `num_steps` parameters. However, not all methods need
discretization along time direction (e.g. method of lines) so this argument
may not be applicable to some implementations.

The workhorse of this method is the `one_step_fn`. For the commonly used
methods, see functions in `math.pde.steppers` module.

The mapping between the arguments of this method and the above
equation are described in the Args section below.

For a simple instructive example of implementation of this method, see
`models.GenericItoProcess.fd_solver_forward`.

Args:
  start_time: Real positive scalar `Tensor`. The start time of the grid.
    Corresponds to time `t0` above.
  end_time: Real scalar `Tensor` smaller than the `start_time` and greater
    than zero. The time to step back to. Corresponds to time `t1` above.
  coord_grid: List of `n` rank 1 real `Tensor`s. `n` is the dimension of the
    domain. The i-th `Tensor` has shape, `[d_i]` where `d_i` is the size of
    the grid along axis `i`. The coordinates of the grid points. Corresponds
    to the spatial grid `G` above.
  values_grid: Real `Tensor` containing the function values at time
    `start_time` which have to be stepped back to time `end_time`. The shape
    of the `Tensor` must broadcast with `[K, d_1, d_2, ..., d_n]`. The first
    axis of size `K` is the values batch dimension and allows multiple
    functions (with potentially different boundary/final conditions) to be
    stepped back simultaneously.
  one_step_fn: The transition kernel. A callable that consumes the following
    arguments by keyword:
      1. 'time': Current time
      2. 'next_time': The next time to step to. For the backwards in time
        evolution, this time will be smaller than the current time.
      3. 'coord_grid': The coordinate grid.
      4. 'values_grid': The values grid.
      5. 'quadratic_coeff': A callable returning the quadratic coefficients
        of the PDE (i.e. `(1/2)D_{ij}(t, x)` above). The callable accepts
        the time and  coordinate grid as keyword arguments and returns a
        `Tensor` with shape that broadcasts with `[dim, dim]`.
      6. 'linear_coeff': A callable returning the linear coefficients of the
        PDE (i.e. `mean_i(t, x)` above). Accepts time and coordinate grid as
        keyword arguments and returns a `Tensor` with shape that broadcasts
        with `[dim]`.
      7. 'constant_coeff': A callable returning the coefficient of the
        linear homogeneous term (i.e. `r(t,x)` above). Same spec as above.
        The `one_step_fn` callable returns a 2-tuple containing the next
        coordinate grid, next values grid.
  boundary_conditions: A list of size `dim` containing boundary conditions.
    The i'th element of the list is a 2-tuple containing the lower and upper
    boundary condition for the boundary along the i`th axis.
  start_step_count: Scalar integer `Tensor`. Initial value for the number of
    time steps performed.
    Default value: 0 (i.e. no previous steps performed).
  num_steps: Positive int scalar `Tensor`. The number of time steps to take
    when moving from `start_time` to `end_time`. Either this argument or the
    `time_step` argument must be supplied (but not both). If num steps is
    `k>=1`, uniform time steps of size `(t0 - t1)/k` are taken to evolve the
    solution from `t0` to `t1`. Corresponds to the `n_steps` parameter
    above.
  time_step: The time step to take. Either this argument or the `num_steps`
    argument must be supplied (but not both). The type of this argument may
    be one of the following (in order of generality): (a) None in which case
      `num_steps` must be supplied. (b) A positive real scalar `Tensor`. The
      maximum time step to take. If the value of this argument is `dt`, then
      the total number of steps taken is N = (t1 - t0) / dt rounded up to
      the nearest integer. The first N-1 steps are of size dt and the last
      step is of size `t1 - t0 - (N-1) * dt`. (c) A callable accepting the
      current time and returning the size of the step to take. The input and
      the output are real scalar `Tensor`s.
  values_transform_fn: An optional callable applied to transform the
    solution values at each time step. The callable is invoked after the
    time step has been performed. The callable should accept the time of the
    grid, the coordinate grid and the values grid and should return the
    values grid. All input arguments to be passed by keyword.
  dtype: The dtype to use.
  name: The name to give to the ops.
    Default value: None which means `solve_forward` is used.
  **kwargs: Additional keyword args:
    (1) pde_solver_fn: Function to solve the PDE that accepts all the above
      arguments by name and returns the same tuple object as required below.
      Defaults to `tff.math.pde.fd_solvers.solve_forward`.

Returns:
  A tuple object containing at least the following attributes:
    final_values_grid: A `Tensor` of same shape and dtype as `values_grid`.
      Contains the final state of the values grid at time `end_time`.
    final_coord_grid: A list of `Tensor`s of the same specification as
      the input `coord_grid`. Final state of the coordinate grid at time
      `end_time`.
    step_count: The total step count (i.e. the sum of the `start_step_count`
      and the number of steps performed in this call.).
    final_time: The final time at which the evolution stopped. This value
      is given by `max(min(end_time, start_time), 0)`.


def eigvals(self, name='eigvals'):
    if not self.is_self_adjoint:
        raise NotImplementedError('Only self-adjoint matrices are supported.')
    with self._name_scope(name):
        return self._eigvals()

Returns the eigenvalues of this linear operator.

If the operator is marked as self-adjoint (via `is_self_adjoint`)
this computation can be more efficient.

Note: This currently only supports self-adjoint operators.

Args:
  name:  A name for this `Op`.

Returns:
  Shape `[B1,...,Bb, N]` `Tensor` of same `dtype` as `self`.


class DepthProImageProcessor(BaseImageProcessor):
    model_input_names = ['pixel_values']

    def __init__(self, do_resize: bool=True, size: Optional[Dict[str, int]]=None, resample: PILImageResampling=PILImageResampling.BILINEAR, do_rescale: bool=True, rescale_factor: Union[int, float]=1 / 255, do_normalize: bool=True, image_mean: Optional[Union[float, List[float]]]=None, image_std: Optional[Union[float, List[float]]]=None, **kwargs):
        super().__init__(**kwargs)
        size = size if size is not None else {'height': 1536, 'width': 1536}
        size = get_size_dict(size)
        self.do_resize = do_resize
        self.do_rescale = do_rescale
        self.do_normalize = do_normalize
        self.size = size
        self.resample = resample
        self.rescale_factor = rescale_factor
        self.image_mean = image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
        self.image_std = image_std if image_std is not None else IMAGENET_STANDARD_STD

    def resize(self, image: np.ndarray, size: Dict[str, int], resample: PILImageResampling=PILImageResampling.BILINEAR, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None, **kwargs) -> np.ndarray:
        """
        Resize an image to `(size["height"], size["width"])`.

        Args:
            image (`np.ndarray`):
                Image to resize.
            size (`Dict[str, int]`):
                Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image.
            resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`):
                `PILImageResampling` filter to use when resizing the image e.g. `PILImageResampling.BILINEAR`.
            data_format (`ChannelDimension` or `str`, *optional*):
                The channel dimension format for the output image. If unset, the channel dimension format of the input
                image is used. Can be one of:
                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
            input_data_format (`ChannelDimension` or `str`, *optional*):
                The channel dimension format for the input image. If unset, the channel dimension format is inferred
                from the input image. Can be one of:
                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.

        Returns:
            `np.ndarray`: The resized images.
        """
        requires_backends(self, 'torch')
        size = get_size_dict(size)
        if 'height' not in size or 'width' not in size:
            raise ValueError(f'The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}')
        output_size = (size['height'], size['width'])
        image_tensor = torch.from_numpy(image).unsqueeze(0)
        resized_image = torch.nn.functional.interpolate(input=image_tensor, size=output_size, mode=pil_torch_interpolation_mapping[resample].value)
        resized_image = resized_image.squeeze(0).numpy()
        return resized_image

    def _validate_input_arguments(self, do_resize: bool, size: Dict[str, int], resample: PILImageResampling, do_rescale: bool, rescale_factor: float, do_normalize: bool, image_mean: Union[float, List[float]], image_std: Union[float, List[float]], data_format: Union[str, ChannelDimension]):
        if do_resize and None in (size, resample):
            raise ValueError('Size and resample must be specified if do_resize is True.')
        if do_rescale and rescale_factor is None:
            raise ValueError('Rescale factor must be specified if do_rescale is True.')
        if do_normalize and None in (image_mean, image_std):
            raise ValueError('Image mean and standard deviation must be specified if do_normalize is True.')

    @filter_out_non_signature_kwargs()
    def preprocess(self, images: ImageInput, do_resize: Optional[bool]=None, size: Optional[Dict[str, int]]=None, resample: Optional[PILImageResampling]=None, do_rescale: Optional[bool]=None, rescale_factor: Optional[float]=None, do_normalize: Optional[bool]=None, image_mean: Optional[Union[float, List[float]]]=None, image_std: Optional[Union[float, List[float]]]=None, return_tensors: Optional[Union[str, TensorType]]=None, data_format: Union[str, ChannelDimension]=ChannelDimension.FIRST, input_data_format: Optional[Union[str, ChannelDimension]]=None):
        """
        Preprocess an image or batch of images.

        Args:
            images (`ImageInput`):
                Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
                passing in images with pixel values between 0 and 1, set `do_rescale=False`.
            do_resize (`bool`, *optional*, defaults to `self.do_resize`):
                Whether to resize the image.
            size (`Dict[str, int]`, *optional*, defaults to `self.size`):
                Dictionary in the format `{"height": h, "width": w}` specifying the size of the output image after
                resizing.
            resample (`PILImageResampling` filter, *optional*, defaults to `self.resample`):
                `PILImageResampling` filter to use if resizing the image e.g. `PILImageResampling.BILINEAR`. Only has
                an effect if `do_resize` is set to `True`.
            do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
                Whether to rescale the image values between [0 - 1].
            rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
                Rescale factor to rescale the image by if `do_rescale` is set to `True`.
            do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
                Whether to normalize the image.
            image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
                Image mean to use if `do_normalize` is set to `True`.
            image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
                Image standard deviation to use if `do_normalize` is set to `True`.
            return_tensors (`str` or `TensorType`, *optional*):
                The type of tensors to return. Can be one of:
                - Unset: Return a list of `np.ndarray`.
                - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
                - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
                - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
                - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
            data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
                The channel dimension format for the output image. Can be one of:
                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
                - Unset: Use the channel dimension format of the input image.
            input_data_format (`ChannelDimension` or `str`, *optional*):
                The channel dimension format for the input image. If unset, the channel dimension format is inferred
                from the input image. Can be one of:
                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
        """
        do_resize = do_resize if do_resize is not None else self.do_resize
        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
        resample = resample if resample is not None else self.resample
        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
        image_mean = image_mean if image_mean is not None else self.image_mean
        image_std = image_std if image_std is not None else self.image_std
        size = size if size is not None else self.size
        images = make_list_of_images(images)
        if not valid_images(images):
            raise ValueError('Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, torch.Tensor, tf.Tensor or jax.ndarray.')
        self._validate_input_arguments(do_resize=do_resize, size=size, resample=resample, do_rescale=do_rescale, rescale_factor=rescale_factor, do_normalize=do_normalize, image_mean=image_mean, image_std=image_std, data_format=data_format)
        images = [to_numpy_array(image) for image in images]
        if is_scaled_image(images[0]) and do_rescale:
            logger.warning_once('It looks like you are trying to rescale already rescaled images. If the input images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again.')
        if input_data_format is None:
            input_data_format = infer_channel_dimension_format(images[0])
        all_images = []
        for image in images:
            if do_rescale:
                image = self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)
            if do_normalize:
                image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
            if do_resize:
                image = to_channel_dimension_format(image, ChannelDimension.FIRST, input_channel_dim=input_data_format)
                image = self.resize(image=image, size=size, resample=resample)
                image = to_channel_dimension_format(image, data_format, input_channel_dim=ChannelDimension.FIRST)
            else:
                image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
            all_images.append(image)
        data = {'pixel_values': all_images}
        return BatchFeature(data=data, tensor_type=return_tensors)

    def post_process_depth_estimation(self, outputs: 'DepthProDepthEstimatorOutput', target_sizes: Optional[Union[TensorType, List[Tuple[int, int]], None]]=None) -> Dict[str, List[TensorType]]:
        """
        Post-processes the raw depth predictions from the model to generate
        final depth predictions which is caliberated using the field of view if provided
        and resized to specified target sizes if provided.

        Args:
            outputs ([`DepthProDepthEstimatorOutput`]):
                Raw outputs of the model.
            target_sizes (`Optional[Union[TensorType, List[Tuple[int, int]], None]]`, *optional*, defaults to `None`):
                Target sizes to resize the depth predictions. Can be a tensor of shape `(batch_size, 2)`
                or a list of tuples `(height, width)` for each image in the batch. If `None`, no resizing
                is performed.

        Returns:
            `List[Dict[str, TensorType]]`: A list of dictionaries of tensors representing the processed depth
            predictions, and field of view (degrees) and focal length (pixels) if `field_of_view` is given in `outputs`.

        Raises:
            `ValueError`:
                If the lengths of `predicted_depths`, `fovs`, or `target_sizes` are mismatched.
        """
        requires_backends(self, 'torch')
        predicted_depth = outputs.predicted_depth
        fov = outputs.field_of_view
        batch_size = len(predicted_depth)
        if target_sizes is not None and batch_size != len(target_sizes):
            raise ValueError('Make sure that you pass in as many fov values as the batch dimension of the predicted depth')
        results = []
        fov = [None] * batch_size if fov is None else fov
        target_sizes = [None] * batch_size if target_sizes is None else target_sizes
        for depth, fov_value, target_size in zip(predicted_depth, fov, target_sizes):
            focal_length = None
            if target_size is not None:
                if fov_value is not None:
                    width = target_size[1]
                    focal_length = 0.5 * width / torch.tan(0.5 * torch.deg2rad(fov_value))
                    depth = depth * width / focal_length
                depth = torch.nn.functional.interpolate(input=depth.unsqueeze(0).unsqueeze(1), size=target_size, mode=pil_torch_interpolation_mapping[self.resample].value).squeeze()
            depth = 1.0 / torch.clamp(depth, min=0.0001, max=10000.0)
            results.append({'predicted_depth': depth, 'field_of_view': fov_value, 'focal_length': focal_length})
        return results

Constructs a DepthPro image processor.

Args:
    do_resize (`bool`, *optional*, defaults to `True`):
        Whether to resize the image's (height, width) dimensions to the specified `(size["height"],
        size["width"])`. Can be overridden by the `do_resize` parameter in the `preprocess` method.
    size (`dict`, *optional*, defaults to `{"height": 1536, "width": 1536}`):
        Size of the output image after resizing. Can be overridden by the `size` parameter in the `preprocess`
        method.
    resample (`PILImageResampling`, *optional*, defaults to `Resampling.BILINEAR`):
        Resampling filter to use if resizing the image. Can be overridden by the `resample` parameter in the
        `preprocess` method.
    do_rescale (`bool`, *optional*, defaults to `True`):
        Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the `do_rescale`
        parameter in the `preprocess` method.
    rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
        Scale factor to use if rescaling the image. Can be overridden by the `rescale_factor` parameter in the
        `preprocess` method.
    do_normalize (`bool`, *optional*, defaults to `True`):
        Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess`
        method.
    image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_MEAN`):
        Mean to use if normalizing the image. This is a float or list of floats the length of the number of
        channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.
    image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_STD`):
        Standard deviation to use if normalizing the image. This is a float or list of floats the length of the
        number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.


class InstructBlipVideoForConditionalGenerationModelOutput(ModelOutput):
    loss: Optional[Tuple[torch.FloatTensor]] = None
    logits: Optional[Tuple[torch.FloatTensor]] = None
    vision_outputs: Optional[torch.FloatTensor] = None
    qformer_outputs: Optional[Tuple[torch.FloatTensor]] = None
    language_model_outputs: Optional[Tuple[torch.FloatTensor]] = None

    def to_tuple(self) -> Tuple[Any]:
        return tuple((self[k] if k not in ['vision_outputs', 'qformer_outputs', 'language_model_outputs'] else getattr(self, k).to_tuple() for k in self.keys()))

Class defining the outputs of [`InstructBlipVideoForConditionalGeneration`].

Args:
    loss (`torch.FloatTensor`, *optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
        Language modeling loss from the language model.
    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
        Prediction scores of the language modeling head of the language model.
    vision_outputs (`BaseModelOutputWithPooling`):
        Outputs of the vision encoder.
    qformer_outputs (`BaseModelOutputWithPoolingAndCrossAttentions`):
        Outputs of the Q-Former (Querying Transformer).
    language_model_outputs (`CausalLMOutputWithPast` or `Seq2SeqLMOutput`):
        Outputs of the language model.


def _delete_blob(self, bucket, blob_name):
    if self._use_blob_generation:
        blob = bucket.get_blob(blob_name, retry=self._storage_client_retry)
        generation = getattr(blob, 'generation', None)
    else:
        generation = None
    try:
        bucket.delete_blob(blob_name, if_generation_match=generation, retry=self._storage_client_retry)
    except NotFound:
        return

Helper method to delete a single blob from GCS.

Args:
  bucket: The GCS bucket object.
  blob_name: The name of the blob to delete under the bucket.


def conv1d_v2(input, filters, stride, padding, data_format='NWC', dilations=None, name=None):
    return conv1d(input, filters, stride, padding, use_cudnn_on_gpu=True, data_format=data_format, name=name, dilations=dilations)

Computes a 1-D convolution given 3-D input and filter tensors.

Given an input tensor of shape
  `batch_shape + [in_width, in_channels]`
if `data_format` is `"NWC"`, or
  `batch_shape + [in_channels, in_width]`
if `data_format` is `"NCW"`,
and a filter / kernel tensor of shape
`[filter_width, in_channels, out_channels]`, this op reshapes
the arguments to pass them to `conv2d` to perform the equivalent
convolution operation.

Internally, this op reshapes the input tensors and invokes `tf.nn.conv2d`.
For example, if `data_format` does not start with `"NC"`, a tensor of shape
  `batch_shape + [in_width, in_channels]`
is reshaped to
  `batch_shape + [1, in_width, in_channels]`,
and the filter is reshaped to
  `[1, filter_width, in_channels, out_channels]`.
The result is then reshaped back to
  `batch_shape + [out_width, out_channels]`
\(where out_width is a function of the stride and padding as in conv2d\) and
returned to the caller.

Args:
  input: A Tensor of rank at least 3. Must be of type `float16`, `float32`, or
    `float64`.
  filters: A Tensor of rank at least 3.  Must have the same type as `input`.
  stride: An int or list of `ints` that has length `1` or `3`.  The number of
    entries by which the filter is moved right at each step.
  padding: 'SAME' or 'VALID'. See
    [here](https://www.tensorflow.org/api_docs/python/tf/nn#notes_on_padding_2)
    for more information.
  data_format: An optional `string` from `"NWC", "NCW"`.  Defaults to `"NWC"`,
    the data is stored in the order of
    `batch_shape + [in_width, in_channels]`.  The `"NCW"` format stores data
    as `batch_shape + [in_channels, in_width]`.
  dilations: An int or list of `ints` that has length `1` or `3` which
    defaults to 1. The dilation factor for each dimension of input. If set to
    k > 1, there will be k-1 skipped cells between each filter element on that
    dimension. Dilations in the batch and depth dimensions must be 1.
  name: A name for the operation (optional).

Returns:
  A `Tensor`.  Has the same type as input.

Raises:
  ValueError: if `data_format` is invalid.


def run_inference(self, batch: Sequence[str], model: _VLLMModelServer, inference_args: Optional[dict[str, Any]]=None) -> Iterable[PredictionResult]:
    return asyncio.run(self._async_run_inference(batch, model, inference_args))

Runs inferences on a batch of text strings.

Args:
  batch: A sequence of examples as text strings.
  model: A _VLLMModelServer containing info for connecting to the server.
  inference_args: Any additional arguments for an inference.

Returns:
  An Iterable of type PredictionResult.


def cross_entropy(self, other, name='cross_entropy'):
    with self._name_scope(name):
        return self._cross_entropy(other)

Computes the (Shannon) cross entropy.

Denote this distribution (`self`) by `P` and the `other` distribution by
`Q`. Assuming `P, Q` are absolutely continuous with respect to
one another and permit densities `p(x) dr(x)` and `q(x) dr(x)`, (Shanon)
cross entropy is defined as:

```none
H[P, Q] = E_p[-log q(X)] = -int_F p(x) log q(x) dr(x)
```

where `F` denotes the support of the random variable `X ~ P`.

Args:
  other: `tfp.distributions.Distribution` instance.
  name: Python `str` prepended to names of ops created by this function.

Returns:
  cross_entropy: `self.dtype` `Tensor` with shape `[B1, ..., Bn]`
    representing `n` different calculations of (Shanon) cross entropy.


def traverse(value: Any, preorder_visitor_fn: Optional[Callable[[KeyPath, Any], bool]]=None, postorder_visitor_fn: Optional[Callable[[KeyPath, Any], bool]]=None, root_path: Optional[KeyPath]=None) -> bool:
    root_path = root_path or KeyPath()

    def no_op_visitor(key, value):
        del key, value
        return True
    if preorder_visitor_fn is None:
        preorder_visitor_fn = no_op_visitor
    if postorder_visitor_fn is None:
        postorder_visitor_fn = no_op_visitor
    if not preorder_visitor_fn(root_path, value):
        return False
    if isinstance(value, dict):
        for k in value.keys():
            if not traverse(value[k], preorder_visitor_fn, postorder_visitor_fn, KeyPath(k, root_path)):
                return False
    elif isinstance(value, list):
        for i, v in enumerate(value):
            if not traverse(v, preorder_visitor_fn, postorder_visitor_fn, KeyPath(i, root_path)):
                return False
    if not postorder_visitor_fn(root_path, value):
        return False
    return True

Traverse a (maybe) hierarchical value.

Example::

  def preorder_visit(path, value):
    print(path)

  tree = {'a': [{'c': [1, 2]}, {'d': {'g': (3, 4)}}], 'b': 'foo'}
  pg.utils.traverse(tree, preorder_visit)

  # Should print:
  # 'a'
  # 'a[0]'
  # 'a[0].c'
  # 'a[0].c[0]'
  # 'a[0].c[1]'
  # 'a[1]'
  # 'a[1].d'
  # 'a[1].d.g'
  # 'b'

Args:
  value: A maybe hierarchical value to traverse.
  preorder_visitor_fn: Preorder visitor function. Function signature is (path,
    value) -> should_continue.
  postorder_visitor_fn: Postorder visitor function. Function signature is
    (path, value) -> should_continue.
  root_path: The key path of the root value.

Returns:
  Whether visitor function returns True on all nodes.


def add_metric(self, value, aggregation=None, name=None):
    if aggregation is not None and aggregation != 'mean':
        raise ValueError('We currently support only `mean` sample-wise metric aggregation. You provided aggregation=`%s`' % aggregation)
    from_metric_obj = hasattr(value, '_metric_obj')
    is_symbolic = tf_utils.is_symbolic_tensor(value)
    in_call_context = base_layer_utils.call_context().in_call
    if name is None and (not from_metric_obj):
        raise ValueError("Please provide a name for your metric like `self.add_metric(tf.reduce_sum(inputs), name='mean_activation', aggregation='mean')`")
    elif from_metric_obj:
        name = value._metric_obj.name
    if in_call_context:
        self._symbolic_add_metric(value, aggregation, name)
    else:
        if not is_symbolic:
            raise ValueError('Expected a symbolic Tensor for the metric value, received: ' + str(value))
        if not getattr(self, '_is_graph_network', False):
            with backend.get_graph().as_default():
                self._symbolic_add_metric(value, aggregation, name)
            return
        if from_metric_obj:
            raise ValueError('Using the result of calling a `Metric` object when calling `add_metric` on a Functional Model is not supported. Please pass the Tensor to monitor directly.')
        self._graph_network_add_metric(value, aggregation, name)

Adds metric tensor to the layer.

Args:
  value: Metric tensor.
  aggregation: Sample-wise metric reduction function. If `aggregation=None`,
    it indicates that the metric tensor provided has been aggregated
    already. eg, `bin_acc = BinaryAccuracy(name='acc')` followed by
    `model.add_metric(bin_acc(y_true, y_pred))`. If aggregation='mean', the
    given metric tensor will be sample-wise reduced using `mean` function.
    eg, `model.add_metric(tf.reduce_sum(outputs), name='output_mean',
    aggregation='mean')`.
  name: String metric name.

Raises:
  ValueError: If `aggregation` is anything other than None or `mean`.


def forward(self, hidden_states: torch.Tensor, position_embeddings: Optional[torch.Tensor]=None, reference_points=None, spatial_shapes=None, spatial_shapes_list=None, level_start_index=None, encoder_hidden_states: Optional[torch.Tensor]=None, encoder_attention_mask: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=False):
    residual = hidden_states
    hidden_states, self_attn_weights = self.self_attn(hidden_states=hidden_states, position_embeddings=position_embeddings, output_attentions=output_attentions)
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    hidden_states = residual + hidden_states
    hidden_states = self.self_attn_layer_norm(hidden_states)
    second_residual = hidden_states
    cross_attn_weights = None
    hidden_states, cross_attn_weights = self.encoder_attn(hidden_states=hidden_states, attention_mask=encoder_attention_mask, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, position_embeddings=position_embeddings, reference_points=reference_points, spatial_shapes=spatial_shapes, spatial_shapes_list=spatial_shapes_list, level_start_index=level_start_index, output_attentions=output_attentions)
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    hidden_states = second_residual + hidden_states
    hidden_states = self.encoder_attn_layer_norm(hidden_states)
    residual = hidden_states
    hidden_states = self.activation_fn(self.fc1(hidden_states))
    hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
    hidden_states = self.fc2(hidden_states)
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    hidden_states = residual + hidden_states
    hidden_states = self.final_layer_norm(hidden_states)
    outputs = (hidden_states,)
    if output_attentions:
        outputs += (self_attn_weights, cross_attn_weights)
    return outputs

Args:
    hidden_states (`torch.FloatTensor`):
        Input to the layer of shape `(seq_len, batch, embed_dim)`.
    position_embeddings (`torch.FloatTensor`, *optional*):
        Position embeddings that are added to the queries and keys in the self-attention layer.
    reference_points (`torch.FloatTensor`, *optional*):
        Reference points.
    spatial_shapes (`torch.LongTensor`, *optional*):
        Spatial shapes.
    level_start_index (`torch.LongTensor`, *optional*):
        Level start index.
    encoder_hidden_states (`torch.FloatTensor`):
        cross attention input to the layer of shape `(seq_len, batch, embed_dim)`
    encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size
        `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative
        values.
    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.


def unique():

    def _apply_fn(dataset):
        return dataset.unique()
    return _apply_fn

Creates a `Dataset` from another `Dataset`, discarding duplicates.

Use this transformation to produce a dataset that contains one instance of
each unique element in the input. For example:

```python
dataset = tf.data.Dataset.from_tensor_slices([1, 37, 2, 37, 2, 1])

# Using `unique()` will drop the duplicate elements.
dataset = dataset.apply(tf.data.experimental.unique())  # ==> { 1, 37, 2 }
```

Returns:
  A `Dataset` transformation function, which can be passed to
  `tf.data.Dataset.apply`.


def write(self, index, value, name=None):
    return self._implementation.write(index, value, name=name)

Write `value` into index `index` of the TensorArray.

Args:
  index: 0-D.  int32 scalar with the index to write to.
  value: N-D.  Tensor of type `dtype`.  The Tensor to write to this index.
  name: A name for the operation (optional).

Returns:
  A new TensorArray object with flow that ensures the write occurs.
  Use this object for all subsequent operations.

Raises:
  ValueError: if there are more writers than specified.


def on_merge(self, to_be_merged, merge_result, context):
    pass

Called when multiple windows are merged.

Args:
  to_be_merged: the set of windows to be merged
  merge_result: the window into which the windows are being merged
  context: a context (e.g. a TriggerContext instance) for managing state
      and setting timers


def apply_tool_use_template(self, conversation: Union[List[Dict[str, str]]], tools: List[Dict], **kwargs) -> Union[str, List[int]]:
    return self.apply_chat_template(conversation, chat_template='tool_use', tools=tools, **kwargs)

Create a Command-R tool-use prompt.

Once rendered, the prompt instructs the model to generate a list of actions to perform on a set of user supplied tools
to help carry out the user's requests.

Conceptually, this works in the same way as `apply_chat_format`, but takes an additional `tools` parameter.

Converts a chat in the form of a list of dictionaries with `"role"` and `"content"` keys and a list of available
tools for the model to use into a prompt string, or a list of token ids.
This method will use the tokenizer's `default_tool_use_template` template specified at the class level.
You can override the default template using the `tool_use_template` kwarg but the quality of your results may decrease.

Args:
    conversation (Union[List[Dict[str, str]]]): A list of dicts
        with "role" and "content" keys, representing the chat history so far.
    tools (List[Dict]): a list of tools to render into the prompt for the model to choose from.
        See an example at the bottom of the docstring.
        The format should be:
           * name (str): The name of the tool to be called. Valid names contain only the characters a-z,
                A-Z, 0-9, _ and must not begin with a digit.
           * description (str): The description of what the tool does, the model uses the description to
                choose when and how to call the function.
           * parameter_definitions (List[Dict]): The input parameters of the tool. Accepts a dictionary
                where the key is the name of the parameter and the value is the parameter spec.
                Valid parameter names contain only the characters a-z, A-Z, 0-9, _ and must not begin with a digit.
                Parameter specs are as follows:
               * description (str): The description of the parameter.
               * type (str): the type of the parameter - most effective for python builtin data types, such as 'str', 'bool'
               * required: boolean: Denotes whether the parameter is always present (required) or not. Defaults to not required.
    add_generation_prompt (bool, *optional*): Whether to end the prompt with the token(s) that indicate
        the start of an assistant message. This is useful when you want to generate a response from the model.
        Note that this argument will be passed to the chat template, and so it must be supported in the
        template for this argument to have any effect.
    tokenize (`bool`, defaults to `True`):
        Whether to tokenize the output. If `False`, the output will be a string.
    padding (`bool`, defaults to `False`):
        Whether to pad sequences to the maximum length. Has no effect if tokenize is `False`.
    truncation (`bool`, defaults to `False`):
        Whether to truncate sequences at the maximum length. Has no effect if tokenize is `False`.
    max_length (`int`, *optional*):
        Maximum length (in tokens) to use for padding or truncation. Has no effect if tokenize is `False`. If
        not specified, the tokenizer's `max_length` attribute will be used as a default.
    return_tensors (`str` or [`~utils.TensorType`], *optional*):
        If set, will return tensors of a particular framework. Has no effect if tokenize is `False`. Acceptable
        values are:
        - `'tf'`: Return TensorFlow `tf.Tensor` objects.
        - `'pt'`: Return PyTorch `torch.Tensor` objects.
        - `'np'`: Return NumPy `np.ndarray` objects.
        - `'jax'`: Return JAX `jnp.ndarray` objects.
    return_dict (`bool`, *optional*, defaults to `False`):
        Whether to return a dictionary with named outputs. Has no effect if tokenize is `False`.
    **tokenizer_kwargs: Additional kwargs to pass to the tokenizer.

Returns:
    `str`: A rendered prompt string.
    or if tokenize=True:
    `List[int]`: A list of token ids representing the tokenized chat so far, including control tokens. This
    output is ready to pass to the model, either directly or via methods like `generate()`.

Examples:

```python
>> tokenizer = CohereTokenizerFast.from_pretrained("CohereForAI/c4ai-command-r-v01")
>> tools = [
    {
        "name": "internet_search",
        "description": "Returns a list of relevant document snippets for a textual query retrieved from the internet",
        "parameter_definitions": {
            "query": {
                "description": "Query to search the internet with",
                "type": "str",
                "required": True
            }
        }
    },
    {
        "name': "directly_answer",
        "description": "Calls a standard (un-augmented) AI chatbot to generate a response given the conversation history",
        "parameter_definitions": {}
    }
]
>> conversation = [
    {"role": "user", "content": "Whats the biggest penguin in the world?"}
]
>> # render the prompt, ready for user to inspect, or for input into the model:
>> prompt = tokenizer.apply_tool_use_template(conversation, tools=tools, tokenize=False, add_generation_prompt=True)
>> print(prompt)
<BOS_TOKEN><|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|># Safety Preamble
The instructions in this section override those in the task description and style guide sections. Don't answer questions that are harmful or immoral.

# System Preamble
## Basic Rules
You are a powerful conversational AI trained by Cohere to help people. You are augmented by a number of tools, and your job is to use and consume the output of these tools to best help the user. You will see a conversation history between yourself and a user, ending with an utterance from the user. You will then see a specific instruction instructing you what kind of response to generate. When you answer the user's requests, you cite your sources in your answers, according to those instructions.

# User Preamble
## Task and Context
You help people answer their questions and other requests interactively. You will be asked a very wide array of requests on all kinds of topics. You will be equipped with a wide range of search engines or similar tools to help you, which you use to research your answer. You should focus on serving the user's needs as best you can, which will be wide-ranging.

## Style Guide
Unless the user asks for a different style of answer, you should answer in full sentences, using proper grammar and spelling.

## Available Tools
Here is a list of tools that you have available to you:

\`\`\`python
def internet_search(query: str) -> List[Dict]:
    """Returns a list of relevant document snippets for a textual query retrieved from the internet

    Args:
        query (str): Query to search the internet with
    """
    pass
\`\`\`

\`\`\`python
def directly_answer() -> List[Dict]:
    """Calls a standard (un-augmented) AI chatbot to generate a response given the conversation history
    """
    pass
\`\`\`<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|USER_TOKEN|>Whats the biggest penguin in the world?<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>Write 'Action:' followed by a json-formatted list of actions that you want to perform in order to produce a good response to the user's last input. You can use any of the supplied tools any number of times, but you should aim to execute the minimum number of necessary actions for the input. You should use the `directly-answer` tool if calling the other tools is unnecessary. The list of actions you want to call should be formatted as a list of json objects, for example:
\`\`\`json
[
    {
        "tool_name": title of the tool in the specification,
        "parameters": a dict of parameters to input into the tool as they are defined in the specs, or {} if it takes no parameters
    }
]\`\`\`<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>
```
>> inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors='pt')
>> outputs = model.generate(inputs, max_new_tokens=128)
>> print(tokenizer.decode(outputs[0]))
Action: ```json
[
    {
        "tool_name": "internet_search",
        "parameters": {
            "query": "biggest penguin in the world"
        }
    }
]
```


def _prepare_images_structure(self, images: ImageInput) -> ImageInput:
    return make_flat_list_of_images(images)

Prepare the images structure for processing.

Args:
    images (`ImageInput`):
        The input images to process.

Returns:
    `ImageInput`: The images with a valid nesting.


def submodules(self):
    return tuple(self._flatten(predicate=_is_module))

Sequence of all sub-modules.

Submodules are modules which are properties of this module, or found as
properties of modules which are properties of this module (and so on).

>>> a = tf.Module()
>>> b = tf.Module()
>>> c = tf.Module()
>>> a.b = b
>>> b.c = c
>>> list(a.submodules) == [b, c]
True
>>> list(b.submodules) == [c]
True
>>> list(c.submodules) == []
True

Returns:
  A sequence of all submodules.


def from_backbone_configs(cls, backbone_config: PretrainedConfig, **kwargs):
    return cls(backbone_config=backbone_config, **kwargs)

Instantiate a [`DFineConfig`] (or a derived class) from a pre-trained backbone model configuration and DETR model
configuration.

    Args:
        backbone_config ([`PretrainedConfig`]):
            The backbone configuration.

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


def add_business_days(self, date_tensor, num_days, roll_convention=constants.BusinessDayConvention.NONE):
    control_deps = []
    if roll_convention == constants.BusinessDayConvention.NONE:
        message = 'Some dates in date_tensor are not business days. Please specify the roll_convention argument.'
        is_bus_day = self.is_business_day(date_tensor)
        control_deps.append(tf.debugging.assert_equal(is_bus_day, True, message=message))
    else:
        date_tensor = self.roll_to_business_day(date_tensor, roll_convention)
    with tf.control_dependencies(control_deps):
        cumul_bus_days_table = self._compute_cumul_bus_days_table()
        cumul_bus_days = self._gather(cumul_bus_days_table, date_tensor.ordinal() - self._ordinal_offset + 1)
        target_cumul_bus_days = cumul_bus_days + num_days
        bus_day_ordinals_table = self._compute_bus_day_ordinals_table()
        ordinals = self._gather(bus_day_ordinals_table, target_cumul_bus_days)
        with tf.control_dependencies(self._assert_ordinals_in_bounds(ordinals)):
            return dt.from_ordinals(ordinals, validate=False)

Adds given number of business days to given dates.

Note that this is different from calling `add_period_and_roll` with
PeriodType.DAY. For example, adding 5 business days to Monday gives the next
Monday (unless there are holidays on this week or next Monday). Adding 5
days and rolling means landing on Saturday and then rolling either to next
Monday or to Friday of the same week, depending on the roll convention.

If any of the dates in `date_tensor` are not business days, they will be
rolled to business days before doing the addition. If `roll_convention` is
`NONE`, and any dates are not business days, an exception is raised.

Args:
  date_tensor: DateTensor of dates to advance from.
  num_days: Tensor of int32 type broadcastable to `date_tensor`.
  roll_convention: BusinessDayConvention. Determines how to roll a date that
    falls on a holiday.

Returns:
  The resulting DateTensor.


class HfArgumentParser(ArgumentParser):
    dataclass_types: Iterable[DataClassType]

    def __init__(self, dataclass_types: Optional[Union[DataClassType, Iterable[DataClassType]]]=None, **kwargs):
        if dataclass_types is None:
            dataclass_types = []
        elif not isinstance(dataclass_types, Iterable):
            dataclass_types = [dataclass_types]
        if 'formatter_class' not in kwargs:
            kwargs['formatter_class'] = ArgumentDefaultsHelpFormatter
        super().__init__(**kwargs)
        if dataclasses.is_dataclass(dataclass_types):
            dataclass_types = [dataclass_types]
        self.dataclass_types = list(dataclass_types)
        for dtype in self.dataclass_types:
            self._add_dataclass_arguments(dtype)

    @staticmethod
    def _parse_dataclass_field(parser: ArgumentParser, field: dataclasses.Field):
        long_options = [f'--{field.name}']
        if '_' in field.name:
            long_options.append(f'--{field.name.replace('_', '-')}')
        kwargs = field.metadata.copy()
        if isinstance(field.type, str):
            raise RuntimeError('Unresolved type detected, which should have been done with the help of `typing.get_type_hints` method by default')
        aliases = kwargs.pop('aliases', [])
        if isinstance(aliases, str):
            aliases = [aliases]
        origin_type = getattr(field.type, '__origin__', field.type)
        if origin_type is Union or (hasattr(types, 'UnionType') and isinstance(origin_type, types.UnionType)):
            if str not in field.type.__args__ and (len(field.type.__args__) != 2 or type(None) not in field.type.__args__):
                raise ValueError(f"Only `Union[X, NoneType]` (i.e., `Optional[X]`) is allowed for `Union` because the argument parser only supports one type per argument. Problem encountered in field '{field.name}'.")
            if type(None) not in field.type.__args__:
                field.type = field.type.__args__[0] if field.type.__args__[1] is str else field.type.__args__[1]
                origin_type = getattr(field.type, '__origin__', field.type)
            elif bool not in field.type.__args__:
                field.type = field.type.__args__[0] if isinstance(None, field.type.__args__[1]) else field.type.__args__[1]
                origin_type = getattr(field.type, '__origin__', field.type)
        bool_kwargs = {}
        if origin_type is Literal or (isinstance(field.type, type) and issubclass(field.type, Enum)):
            if origin_type is Literal:
                kwargs['choices'] = field.type.__args__
            else:
                kwargs['choices'] = [x.value for x in field.type]
            kwargs['type'] = make_choice_type_function(kwargs['choices'])
            if field.default is not dataclasses.MISSING:
                kwargs['default'] = field.default
            else:
                kwargs['required'] = True
        elif field.type is bool or field.type == Optional[bool]:
            bool_kwargs = copy(kwargs)
            kwargs['type'] = string_to_bool
            if field.type is bool or (field.default is not None and field.default is not dataclasses.MISSING):
                default = False if field.default is dataclasses.MISSING else field.default
                kwargs['default'] = default
                kwargs['nargs'] = '?'
                kwargs['const'] = True
        elif isclass(origin_type) and issubclass(origin_type, list):
            kwargs['type'] = field.type.__args__[0]
            kwargs['nargs'] = '+'
            if field.default_factory is not dataclasses.MISSING:
                kwargs['default'] = field.default_factory()
            elif field.default is dataclasses.MISSING:
                kwargs['required'] = True
        else:
            kwargs['type'] = field.type
            if field.default is not dataclasses.MISSING:
                kwargs['default'] = field.default
            elif field.default_factory is not dataclasses.MISSING:
                kwargs['default'] = field.default_factory()
            else:
                kwargs['required'] = True
        parser.add_argument(*long_options, *aliases, **kwargs)
        if field.default is True and (field.type is bool or field.type == Optional[bool]):
            bool_kwargs['default'] = False
            parser.add_argument(f'--no_{field.name}', f'--no-{field.name.replace('_', '-')}', action='store_false', dest=field.name, **bool_kwargs)

    def _add_dataclass_arguments(self, dtype: DataClassType):
        if hasattr(dtype, '_argument_group_name'):
            parser = self.add_argument_group(dtype._argument_group_name)
        else:
            parser = self
        try:
            type_hints: dict[str, type] = get_type_hints(dtype)
        except NameError:
            raise RuntimeError(f'Type resolution failed for {dtype}. Try declaring the class in global scope or removing line of `from __future__ import annotations` which opts in Postponed Evaluation of Annotations (PEP 563)')
        except TypeError as ex:
            if sys.version_info[:2] < (3, 10) and 'unsupported operand type(s) for |' in str(ex):
                python_version = '.'.join(map(str, sys.version_info[:3]))
                raise RuntimeError(f'Type resolution failed for {dtype} on Python {python_version}. Try removing line of `from __future__ import annotations` which opts in union types as `X | Y` (PEP 604) via Postponed Evaluation of Annotations (PEP 563). To support Python versions that lower than 3.10, you need to use `typing.Union[X, Y]` instead of `X | Y` and `typing.Optional[X]` instead of `X | None`.') from ex
            raise
        for field in dataclasses.fields(dtype):
            if not field.init:
                continue
            field.type = type_hints[field.name]
            self._parse_dataclass_field(parser, field)

    def parse_args_into_dataclasses(self, args=None, return_remaining_strings=False, look_for_args_file=True, args_filename=None, args_file_flag=None) -> tuple[DataClass, ...]:
        """
        Parse command-line args into instances of the specified dataclass types.

        This relies on argparse's `ArgumentParser.parse_known_args`. See the doc at:
        docs.python.org/3.7/library/argparse.html#argparse.ArgumentParser.parse_args

        Args:
            args:
                List of strings to parse. The default is taken from sys.argv. (same as argparse.ArgumentParser)
            return_remaining_strings:
                If true, also return a list of remaining argument strings.
            look_for_args_file:
                If true, will look for a ".args" file with the same base name as the entry point script for this
                process, and will append its potential content to the command line args.
            args_filename:
                If not None, will uses this file instead of the ".args" file specified in the previous argument.
            args_file_flag:
                If not None, will look for a file in the command-line args specified with this flag. The flag can be
                specified multiple times and precedence is determined by the order (last one wins).

        Returns:
            Tuple consisting of:

                - the dataclass instances in the same order as they were passed to the initializer.abspath
                - if applicable, an additional namespace for more (non-dataclass backed) arguments added to the parser
                  after initialization.
                - The potential list of remaining argument strings. (same as argparse.ArgumentParser.parse_known_args)
        """
        if args_file_flag or args_filename or (look_for_args_file and len(sys.argv)):
            args_files = []
            if args_filename:
                args_files.append(Path(args_filename))
            elif look_for_args_file and len(sys.argv):
                args_files.append(Path(sys.argv[0]).with_suffix('.args'))
            if args_file_flag:
                args_file_parser = ArgumentParser()
                args_file_parser.add_argument(args_file_flag, type=str, action='append')
                cfg, args = args_file_parser.parse_known_args(args=args)
                cmd_args_file_paths = vars(cfg).get(args_file_flag.lstrip('-'), None)
                if cmd_args_file_paths:
                    args_files.extend([Path(p) for p in cmd_args_file_paths])
            file_args = []
            for args_file in args_files:
                if args_file.exists():
                    file_args += args_file.read_text().split()
            args = file_args + args if args is not None else file_args + sys.argv[1:]
        namespace, remaining_args = self.parse_known_args(args=args)
        outputs = []
        for dtype in self.dataclass_types:
            keys = {f.name for f in dataclasses.fields(dtype) if f.init}
            inputs = {k: v for k, v in vars(namespace).items() if k in keys}
            for k in keys:
                delattr(namespace, k)
            obj = dtype(**inputs)
            outputs.append(obj)
        if len(namespace.__dict__) > 0:
            outputs.append(namespace)
        if return_remaining_strings:
            return (*outputs, remaining_args)
        else:
            if remaining_args:
                raise ValueError(f'Some specified arguments are not used by the HfArgumentParser: {remaining_args}')
            return (*outputs,)

    def parse_dict(self, args: dict[str, Any], allow_extra_keys: bool=False) -> tuple[DataClass, ...]:
        """
        Alternative helper method that does not use `argparse` at all, instead uses a dict and populating the dataclass
        types.

        Args:
            args (`dict`):
                dict containing config values
            allow_extra_keys (`bool`, *optional*, defaults to `False`):
                Defaults to False. If False, will raise an exception if the dict contains keys that are not parsed.

        Returns:
            Tuple consisting of:

                - the dataclass instances in the same order as they were passed to the initializer.
        """
        unused_keys = set(args.keys())
        outputs = []
        for dtype in self.dataclass_types:
            keys = {f.name for f in dataclasses.fields(dtype) if f.init}
            inputs = {k: v for k, v in args.items() if k in keys}
            unused_keys.difference_update(inputs.keys())
            obj = dtype(**inputs)
            outputs.append(obj)
        if not allow_extra_keys and unused_keys:
            raise ValueError(f'Some keys are not used by the HfArgumentParser: {sorted(unused_keys)}')
        return tuple(outputs)

    def parse_json_file(self, json_file: Union[str, os.PathLike], allow_extra_keys: bool=False) -> tuple[DataClass, ...]:
        """
        Alternative helper method that does not use `argparse` at all, instead loading a json file and populating the
        dataclass types.

        Args:
            json_file (`str` or `os.PathLike`):
                File name of the json file to parse
            allow_extra_keys (`bool`, *optional*, defaults to `False`):
                Defaults to False. If False, will raise an exception if the json file contains keys that are not
                parsed.

        Returns:
            Tuple consisting of:

                - the dataclass instances in the same order as they were passed to the initializer.
        """
        with open(Path(json_file), encoding='utf-8') as open_json_file:
            data = json.loads(open_json_file.read())
        outputs = self.parse_dict(data, allow_extra_keys=allow_extra_keys)
        return tuple(outputs)

    def parse_yaml_file(self, yaml_file: Union[str, os.PathLike], allow_extra_keys: bool=False) -> tuple[DataClass, ...]:
        """
        Alternative helper method that does not use `argparse` at all, instead loading a yaml file and populating the
        dataclass types.

        Args:
            yaml_file (`str` or `os.PathLike`):
                File name of the yaml file to parse
            allow_extra_keys (`bool`, *optional*, defaults to `False`):
                Defaults to False. If False, will raise an exception if the json file contains keys that are not
                parsed.

        Returns:
            Tuple consisting of:

                - the dataclass instances in the same order as they were passed to the initializer.
        """
        outputs = self.parse_dict(yaml.safe_load(Path(yaml_file).read_text()), allow_extra_keys=allow_extra_keys)
        return tuple(outputs)

This subclass of `argparse.ArgumentParser` uses type hints on dataclasses to generate arguments.

The class is designed to play well with the native argparse. In particular, you can add more (non-dataclass backed)
arguments to the parser after initialization and you'll get the output back after parsing as an additional
namespace. Optional: To create sub argument groups use the `_argument_group_name` attribute in the dataclass.

Args:
    dataclass_types (`DataClassType` or `Iterable[DataClassType]`, *optional*):
        Dataclass type, or list of dataclass types for which we will "fill" instances with the parsed args.
    kwargs (`Dict[str, Any]`, *optional*):
        Passed to `argparse.ArgumentParser()` in the regular way.


def __init__(self, cluster_resolver=None, communication_options=None, *, mesh=None):
    self._validate_init_args(mesh, cluster_resolver)
    if not mesh:
        if not cluster_resolver:
            cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
        dtensor_env_var = _parse_dtensor_env_var_from_cluster_resolver(cluster_resolver)
        _config_dtensor_env_var(dtensor_env_var)
        mesh = _build_distributed_mesh(dtensor_util.DEFAULT_BATCH_MESH_DIM_NAME)
    extended = dtensor_strategy_extended.DTensorStrategyExtended(container_strategy=self, mesh=mesh)
    super().__init__(extended)
    self._mesh = mesh
    self._cluster_resolver = cluster_resolver

Creates the strategy.

Args:
  cluster_resolver: optional
    `tf.distribute.cluster_resolver.ClusterResolver`. In case neither `mesh`
    nor `cluster_resolver` are provided,
    `tf.distribute.cluster_resolver.TFConfigClusterResolver` is used.
  communication_options: currently ignore.
  mesh: optional Dtensor global mesh for the computation. Note that either
    `mesh` or the `cluster_resolver` should be provided. and not both.


def transform(self, X):
    sklearn.base.check_is_fitted(self)
    X = _validate_data(self, X, reset=False)
    return self.model_.predict(X)

Transform the data.

Args:
    X: array-like, shape=(n_samples, n_features)
        The input samples.

Returns:
    X_transformed: array-like, shape=(n_samples, n_features)
        The transformed data.


def stop_on_exception(self):
    return self._coord.stop_on_exception()

Context handler to stop the supervisor when an exception is raised.

See `Coordinator.stop_on_exception()`.

Returns:
  A context handler.


def build_inputs_with_special_tokens(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None) -> List[int]:
    if token_ids_1 is None:
        return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
    cls = [self.cls_token_id]
    sep = [self.sep_token_id]
    return cls + token_ids_0 + sep + token_ids_1 + sep

Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
adding special tokens. A BERT sequence has the following format:

- single sequence: `[CLS] X [SEP]`
- pair of sequences: `[CLS] A [SEP] B [SEP]`

Args:
    token_ids_0 (`List[int]`):
        List of IDs to which the special tokens will be added.
    token_ids_1 (`List[int]`, *optional*):
        Optional second list of IDs for sequence pairs.

Returns:
    `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.


def pack_sequence_as(structure, flat_sequence):
    return tree_impl.pack_sequence_as(structure, flat_sequence)

Returns a given flattened sequence packed into a given structure.

If `structure` is an atom, `flat_sequence` must be a single-item list; in
this case the return value is `flat_sequence[0]`.

If `structure` is or contains a dict instance, the keys will be sorted to
pack the flat sequence in deterministic order. However, instances of
`collections.OrderedDict` are handled differently: their sequence order is
used instead of the sorted keys. The same convention is followed in
`flatten`. This correctly repacks dicts and `OrderedDicts` after they have
been flattened, or vice-versa.

Dictionaries with non-sortable keys are not supported.

Examples:

>>> structure = {"key3": "", "key1": "", "key2": ""}
>>> flat_sequence = ["value1", "value2", "value3"]
>>> keras.tree.pack_sequence_as(structure, flat_sequence)
{"key3": "value3", "key1": "value1", "key2": "value2"}

>>> structure = (("a", "b"), ("c", "d", "e"), "f")
>>> flat_sequence = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
>>> keras.tree.pack_sequence_as(structure, flat_sequence)
((1.0, 2.0), (3.0, 4.0, 5.0), 6.0)

>>> structure = {"key3": {"c": ("alpha", "beta"), "a": ("gamma")},
... "key1": {"e": "val1", "d": "val2"}}
>>> flat_sequence = ["val2", "val1", 3.0, 1.0, 2.0]
>>> keras.tree.pack_sequence_as(structure, flat_sequence)
{'key3': {'c': (1.0, 2.0), 'a': 3.0}, 'key1': {'e': 'val1', 'd': 'val2'}}

>>> structure = ["a"]
>>> flat_sequence = [np.array([[1, 2], [3, 4]])]
>>> keras.tree.pack_sequence_as(structure, flat_sequence)
[array([[1, 2],
   [3, 4]])]

>>> structure = ["a"]
>>> flat_sequence = [keras.ops.ones([2, 2])]
>>> keras.tree.pack_sequence_as(structure, flat_sequence)
[array([[1., 1.],
   [1., 1.]]]

Args:
    structure: Arbitrarily nested structure.
    flat_sequence: Flat sequence to pack.

Returns:
    `flat_sequence` converted to have the same recursive structure as
    `structure`.

Raises:
    TypeError: If `flat_sequence` is not iterable.
    ValueError: If `flat_sequence` cannot be repacked as `structure`; for
        instance, if `flat_sequence` has too few or too many elements.


def score_one(self, x: beam.Row) -> Optional[float]:
    if len(x.__dict__) != 1:
        raise ValueError('RobustZScore.score_one expected univariate input, but got %s', str(x))
    v = next(iter(x))
    if v is None or math.isnan(v):
        return None
    median = self._mad_tracker.get_median()
    mad = self._mad_tracker.get()
    if math.isnan(mad) or math.isnan(median):
        return float('NaN')
    if abs(mad) < EPSILON:
        return 0.0
    return abs(RobustZScore.SCALE_FACTOR * (v - median) / mad)

Scores a data point using the Robust Z-Score.

Args:
  x: A `beam.Row` containing a single numerical value.

Returns:
  float | None: The Robust Z-Score.


def get_summary_description(node_def):
    if node_def.op != 'TensorSummary':
        raise ValueError("Can't get_summary_description on %s" % node_def.op)
    description_str = _compat.as_str_any(node_def.attr['description'].s)
    summary_description = SummaryDescription()
    _json_format.Parse(description_str, summary_description)
    return summary_description

Given a TensorSummary node_def, retrieve its SummaryDescription.

When a Summary op is instantiated, a SummaryDescription of associated
metadata is stored in its NodeDef. This method retrieves the description.

Args:
  node_def: the node_def_pb2.NodeDef of a TensorSummary op

Returns:
  a summary_pb2.SummaryDescription

Raises:
  ValueError: if the node is not a summary op.

@compatibility(eager)
Not compatible with eager execution. To write TensorBoard
summaries under eager execution, use `tf.contrib.summary` instead.
@end_compatibility


def get_module_action(self, module):
    f = module.full_path
    if f in self.filenames:
        action = Action.CHECK
        report = logging.warning
    else:
        action = Action.INFER
        report = logging.info
    if not module.name.startswith('pytype_extensions.') and module.kind in ('Builtin', 'System'):
        action = Action.GENERATE_DEFAULT
        report('%s: %s module %s', action, module.kind, module.name)
    return action

Get the action for the given module.

Args:
  module: A module_utils.Module object.

Returns:
  An Action object, or None for a non-Python file.