def setup_mock_socket_file(self, mock_create_connection, resp=MOCK_RESP):
    fake_file = self.MockSocketFile(resp)
    fake_conn = mock.MagicMock()
    fake_conn.makefile.return_value = fake_file
    mock_create_connection.return_value = fake_conn
    return fake_file

Sets up a fake socket file from the mock connection.

Args:
        mock_create_connection: The mock method for creating a method.
        resp: (str) response to give. MOCK_RESP by default.

Returns:
        The mock file that will be injected into the code.


def rescale(self, image: 'torch.Tensor', scale: float, offset: Optional[bool]=True, **kwargs) -> 'torch.Tensor':
    rescaled_image = image * scale
    if offset:
        rescaled_image -= 1
    return rescaled_image

Rescale an image by a scale factor.

    If `offset` is `True`, the image has its values rescaled by `scale` and then offset by 1. If `scale` is
    1/127.5, the image is rescaled between [-1, 1].
        image = image * scale - 1

    If `offset` is `False`, and `scale` is 1/255, the image is rescaled between [0, 1].
        image = image * scale

Args:
        image (`torch.Tensor`):
            Image to rescale.
        scale (`float`):
            The scaling factor to rescale pixel values by.
        offset (`bool`, *optional*):
            Whether to scale the image in both negative and positive directions.

Returns:
        `torch.Tensor`: The rescaled image.


def convert_to_sliceable(arrays, target_backend=None):

    def convert_single_array(x):
        if x is None:
            return x
        if isinstance(x, np.ndarray) and str(x.dtype) == 'object' and (backend.backend() == 'tensorflow') and all((isinstance(e, str) for e in x)):
            x = tf.convert_to_tensor(x, dtype='string')
        if isinstance(x, np.ndarray):
            sliceable_class = NumpySliceable
        elif data_adapter_utils.is_tensorflow_tensor(x):
            if data_adapter_utils.is_tensorflow_ragged(x):
                sliceable_class = TensorflowRaggedSliceable
            elif data_adapter_utils.is_tensorflow_sparse(x):
                sliceable_class = TensorflowSparseSliceable
            else:
                sliceable_class = TensorflowSliceable
        elif data_adapter_utils.is_jax_array(x):
            if data_adapter_utils.is_jax_sparse(x):
                sliceable_class = JaxSparseSliceable
            else:
                x = np.asarray(x)
                sliceable_class = NumpySliceable
        elif data_adapter_utils.is_torch_tensor(x):
            sliceable_class = TorchSliceable
        elif pandas is not None and isinstance(x, pandas.DataFrame):
            sliceable_class = PandasDataFrameSliceable
        elif pandas is not None and isinstance(x, pandas.Series):
            sliceable_class = PandasSeriesSliceable
        elif data_adapter_utils.is_scipy_sparse(x):
            sliceable_class = ScipySparseSliceable
        elif hasattr(x, '__array__'):
            x = np.asarray(x)
            sliceable_class = NumpySliceable
        else:
            raise ValueError(f'Expected a NumPy array, tf.Tensor, tf.RaggedTensor, tf.SparseTensor, jax.np.ndarray, jax.experimental.sparse.JAXSparse, torch.Tensor, Pandas Dataframe, or Pandas Series. Received invalid input: {x} (of type {type(x)})')

        def is_non_floatx_float(dtype):
            return dtype is not object and backend.is_float_dtype(dtype) and (not backend.standardize_dtype(dtype) == backend.floatx())
        cast_dtype = None
        if pandas is not None and isinstance(x, pandas.DataFrame):
            if any((is_non_floatx_float(d) for d in x.dtypes.values)):
                cast_dtype = backend.floatx()
        elif is_non_floatx_float(x.dtype):
            cast_dtype = backend.floatx()
        if cast_dtype is not None:
            x = sliceable_class.cast(x, cast_dtype)
        if target_backend is None:
            return sliceable_class(x)
        if target_backend == 'tensorflow':
            return sliceable_class.convert_to_tf_dataset_compatible(x)
        if target_backend == 'jax' and sliceable_class in (TensorflowSliceable, TorchSliceable):
            x = np.asarray(x)
            sliceable_class = NumpySliceable
        return sliceable_class(x)
    return tree.map_structure(convert_single_array, arrays)

Convert a structure of arrays into `Sliceable` instances

Args:
        arrays: the arrays to convert.
        target_backend: the target backend for the output:
            - `None` indicates that `arrays` will be wrapped into `Sliceable`s
            as-is without using a different representation. This is used by
            `train_validation_split()`.
            - `tensorflow` indicates that
            `Sliceable.convert_to_tf_dataset_compatible` will be called. The
            returned structure therefore contains arrays, not `Sliceable`s.
            - `numpy`, `jax` or `torch` indices that the arrays will eventually
            be converted to this backend type after slicing. In this case,
            the intermediary `Sliceable`s may use a different representation
            from the input `arrays` for better performance.
    Returns: the same structure with `Sliceable` instances or arrays.


def before(self: EventSetOrNode, timestamp: Union[int, float, datetime]) -> EventSetOrNode:
    from temporian.core.operators.filter import before
    return before(self, timestamp=timestamp)

Filters events [`EventSet`][temporian.EventSet] that happened before
    a particular timestamp.

    The timestamp can be a datetime if the EventSet's timestamps are unix
    timestamps.

    The comparison is strict, meaning that the obtained timestamps would be
    less than (`<`) the provided timestamp.

    This operation is equivalent to:
    `input.filter(input.timestamps() < timestamp)`

    Usage example:
        ```python
        >>> a = tp.event_set(
        ...     timestamps=[0, 1, 5, 6],
        ...     features={"f1": [0, 10, 50, 60]},
        ... )

        >>> a.before(5)
        indexes: []
        features: [('f1', int64)]
        events:
            (2 events):
                timestamps: [0. 1.]
                'f1': [ 0 10]
        ...

        >>> from datetime import datetime
        >>> a = tp.event_set(
        ...     timestamps=[datetime(2022, 1, 1), datetime(2022, 1, 2)],
        ...     features={"f1": [1, 2]},
        ... )

        >>> a.before(datetime(2022, 1, 1, 12))
        indexes: []
        features: [('f1', int64)]
        events:
            (1 events):
                timestamps: ['2022-01-01T00:00:00']
                'f1': [1]
        ...

        ```

Args:
        timestamp: EventSet with a single boolean feature.

Returns:
        Filtered EventSet.


def compute_gradients(self, loss, var_list=None, **kwargs):
    num_shards = tpu_function.get_tpu_context().number_of_shards
    if num_shards is None:
        logging.warning('CrossShardOptimizer should be used within a tpu_shard_context, but got unset number_of_shards. Assuming 1.')
        num_shards = 1
    subgroup_size = self._verify_and_get_subgroup_size(self._group_assignment, num_shards)
    if num_shards > 1 and self._reduction == losses.Reduction.MEAN:
        if self._group_assignment:
            scale = 1.0 / subgroup_size
        else:
            scale = 1.0 / num_shards
        loss *= scale
    return self._opt.compute_gradients(loss, var_list=var_list, **kwargs)

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

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

    When the CrossShardOptimizer is constructed with
    `reduction == losses.Reduction.MEAN` (default), this function scales the
    loss by `1.0 / num_shards` before computing the gradients. Assuming the
    optimizer uses the default implementation of `compute_gradients()`, the
    gradients of the scaled loss are scaled by `1.0 / num_shards` compared to
    the gradients of the original loss. This scaling factor is important because
    `apply_gradients()` sums gradients across shards, rather than averaging
    them. However, the scaling factor must be taken into account when clipping
    the norm of the gradients or performing other postprocessing.

Args:
        loss: A Tensor containing the value to minimize.
        var_list: Optional list or tuple of `tf.Variable` to update to minimize
        `loss`.  Defaults to the list of variables collected in the graph
        under the key `GraphKey.TRAINABLE_VARIABLES`.
        **kwargs: Keyword arguments for compute_gradients().

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

Raises:
        ValueError: If not within a tpu_shard_context or group_assignment is
        invalid.


def get_associated_prs(api: github_api.GitHubAPI, commit_hashes: Sequence[str]) -> Generator[int, None, None]:
    regex = re.compile('PR #(\\d+)')
    for commit_hash in commit_hashes:
        response = api.get_commit('openxla/xla', commit_hash)
        message = response['commit']['message']
        if (maybe_match := regex.match(message)):
            pr_number = maybe_match.group(1)
            print(f'Found PR #{pr_number} associated with commit hash {commit_hash}')
            yield int(pr_number)
    print(f"Didn't find any PRs associated with commit hashes: {commit_hashes}")

Finds PRs associated with commits.

Args:
        api: GitHubAPI object which will be used to make requests
        commit_hashes: A sequence of SHAs which may have PRs associated with them

Yields:
        Associated pairs of (PR number, SHA), both as strings


def object_graph_key_mapping(checkpoint_path):
    reader = py_checkpoint_reader.NewCheckpointReader(checkpoint_path)
    object_graph_string = reader.get_tensor(trackable.OBJECT_GRAPH_PROTO_KEY)
    object_graph_proto = trackable_object_graph_pb2.TrackableObjectGraph()
    object_graph_proto.ParseFromString(object_graph_string)
    names_to_keys = {}
    for node in object_graph_proto.nodes:
        for attribute in node.attributes:
            names_to_keys[attribute.full_name] = attribute.checkpoint_key
    return names_to_keys

Return name to key mappings from the checkpoint.

Args:
        checkpoint_path: string, path to object-based checkpoint

Returns:
        Dictionary mapping tensor names to checkpoint keys.


def _gather_implementation(self, per_replica_value, destinations, axis, options):
    raise NotImplementedError('_gather method must be implemented in descendants.')

Implementation of `gather` method of `tf.distribute.CrossDeviceOps`.

    Overriding this method is useful for subclass implementers.

Args:
        per_replica_value: a `tf.distribute.DistributedValues`, or a `tf.Tensor`
        like object.
        destinations: a `tf.distribute.DistributedValues`, a `tf.Variable`, a
        `tf.Tensor` alike object, or a device string. It specifies the devices
        to gather to. To perform an all-gather, pass the same to `value` and
        `destinations`. Note that if it's a `tf.Variable`, the value is gathered
        to the devices of that variable, this method doesn't update the
        variable.
        axis: specifies the dimension to gather along within each replica's
        tensor.
        options: a `tf.distribute.experimental.CommunicationOptions`. See
        `tf.distribute.experimental.CommunicationOptions` for details.

Returns:
        A `tf.Tensor` or `tf.distribute.DistributedValues`.

Raises:
        ValueError: if per_replica_value can't be converted to a
        `tf.distribute.DistributedValues` or if destinations is not a string,
        `tf.Variable` or `tf.distribute.DistributedValues`.


def binary_crossentropy(target, output, from_logits=False):
    if any_symbolic_tensors((target, output)):
        return BinaryCrossentropy(from_logits=from_logits).symbolic_call(target, output)
    return backend.nn.binary_crossentropy(target, output, from_logits=from_logits)

Computes binary cross-entropy loss between target and output tensor.

    The binary cross-entropy loss is commonly used in binary
    classification tasks where each input sample belongs to one
    of the two classes. It measures the dissimilarity between the
    target and output probabilities or logits.

Args:
        target: The target tensor representing the true binary labels.
            Its shape should match the shape of the `output` tensor.
        output: The output tensor representing the predicted probabilities
            or logits. Its shape should match the shape of the
            `target` tensor.
        from_logits: (optional) Whether `output` is a tensor of logits or
            probabilities.
            Set it to `True` if `output` represents logits; otherwise,
            set it to `False` if `output` represents probabilities.
            Defaults to `False`.

Returns:
        Integer tensor: The computed binary cross-entropy loss between
        `target` and `output`.

Example:
    >>> target = keras.ops.convert_to_tensor([0, 1, 1, 0])
    >>> output = keras.ops.convert_to_tensor([0.1, 0.9, 0.8, 0.2])
    >>> binary_crossentropy(target, output)
    array([0.10536054 0.10536054 0.22314355 0.22314355],
        shape=(4,), dtype=float32)


def _run_calibration(saved_model_path: str, signature_keys: Sequence[str], tags: Collection[str], force_graph_mode_calibration: bool, representative_dataset_file_map: Mapping[str, quantization_options_pb2.RepresentativeDatasetFile]) -> bool:
    repr_dataset_map = rd.TfRecordRepresentativeDatasetLoader(representative_dataset_file_map).load()
    _run_graph_for_calibration(saved_model_path, signature_keys, tags, repr_dataset_map, force_graph_mode_calibration)
    return True

Runs calibration and adds calibration statistics to exported model.

Args:
        saved_model_path: Path to the SavedModel to run calibration.
        signature_keys: List of signature keys corresponding to SignatureDefs to run
        calibration on.
        tags: A set of tags that identify the MetaGraphDef.
        force_graph_mode_calibration: If True, runs the calibration in graph mode.
        representative_dataset_file_map: Signature key ->
        `RepresentativeDatasetFile` mapping for running the calibration step. Each
        dataset file stores the representative dataset for the function matching
        the signature key.

Returns:
        `True` upon successfully running calibration.


def create_slot_with_initializer(primary, initializer, shape, dtype, name, colocate_with_primary=True, *, copy_xla_sharding=False):
    validate_shape = shape.is_fully_defined()
    if isinstance(primary, variables.Variable):
        prefix = primary._shared_name
    else:
        prefix = primary.op.name
    with variable_scope.variable_scope(None, prefix + '/' + name):
        if colocate_with_primary:
            distribution_strategy = distribute_lib.get_strategy()
            with distribution_strategy.extended.colocate_vars_with(primary):
                return _create_slot_var(primary, initializer, '', validate_shape, shape, dtype, copy_xla_sharding=copy_xla_sharding)
        else:
            return _create_slot_var(primary, initializer, '', validate_shape, shape, dtype, copy_xla_sharding=copy_xla_sharding)

Creates a slot initialized using an `Initializer`.

    The type of the slot is determined by the given value.

Args:
        primary: The primary `Variable` or `Tensor`.
        initializer: An `Initializer`.  The initial value of the slot.
        shape: Shape of the initial value of the slot.
        dtype: Type of the value of the slot.
        name: Name to use for the slot variable.
        colocate_with_primary: Boolean.  If True the slot is located
        on the same device as `primary`.
        copy_xla_sharding: Boolean. If True also copies XLA sharding
        from primary.

Returns:
        A `Variable` object.


def generate_row(fake: Faker, config: Config) -> Row:
    row: Row = {}
    for column in config:
        row[column.bq_name] = column.value(fake)
    return row

Generates a Row of Faker data, conforming to the config.

Args:
        * fake: Faker instance
        * config: List of Columns

Returns:
        * Row of Faker data


def generate_examples(options):
    _prepare_dir(options)
    out = options.zip_to_output
    if options.multi_gen_state:
        test_name = options.multi_gen_state.test_name
    else:
        test_name = re.sub('(_(|with-flex|forward-compat|edgetpu|mlir-quant))?(_xnnpack)?\\.zip$', '', out, count=1)
    test_function_name = 'make_%s_tests' % test_name
    test_function = get_test_function(test_function_name)
    if test_function is None:
        raise RuntimeError("Can't find a test function to create %r. Tried %r" % (out, test_function_name))
    if options.make_forward_compat_test:
        future_date = datetime.date.today() + datetime.timedelta(days=30)
        with tf.compat.forward_compatibility_horizon(future_date.year, future_date.month, future_date.day):
            test_function(options)
    else:
        test_function(options)

Generate examples for a test set.

Args:
        options: Options containing information to generate examples.

Raises:
        RuntimeError: if the test function cannot be found.


def saved_model_format_scope(value, **kwargs):
    previous_format = _thread_local_data.saved_model_format
    previous_kwargs = _thread_local_data.save_kwargs
    try:
        _thread_local_data.saved_model_format = value
        _thread_local_data.save_kwargs = kwargs
        yield
    finally:
        _thread_local_data.saved_model_format = previous_format
        _thread_local_data.save_kwargs = previous_kwargs

Provides a scope within which the savde model format to test is `value`.

    The saved model format gets restored to its original value upon exiting the
    scope.

Args:
        value: saved model format value
        **kwargs: optional kwargs to pass to the save function.

Yields:
        The provided value.


def post_process(self, outputs, target_sizes):
    logging.warning_once('`post_process` is deprecated and will be removed in v5 of Transformers, please use `post_process_object_detection` instead, with `threshold=0.` for equivalent results.')
    out_logits, out_bbox = (outputs.logits, outputs.pred_boxes)
    if len(out_logits) != len(target_sizes):
        raise ValueError('Make sure that you pass in as many target sizes as the batch dimension of the logits')
    if target_sizes.shape[1] != 2:
        raise ValueError('Each element of target_sizes must contain the size (h, w) of each image of the batch')
    prob = out_logits.sigmoid()
    topk_values, topk_indexes = torch.topk(prob.view(out_logits.shape[0], -1), 300, dim=1)
    scores = topk_values
    topk_boxes = torch.div(topk_indexes, out_logits.shape[2], rounding_mode='floor')
    labels = topk_indexes % out_logits.shape[2]
    boxes = center_to_corners_format(out_bbox)
    boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4))
    img_h, img_w = target_sizes.unbind(1)
    scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1)
    boxes = boxes * scale_fct[:, None, :]
    results = [{'scores': s, 'labels': l, 'boxes': b} for s, l, b in zip(scores, labels, boxes)]
    return results

Converts the output of [`ConditionalDetrForObjectDetection`] into the format expected by the Pascal VOC format (xmin, ymin, xmax, ymax).
    Only supports PyTorch.

Args:
        outputs ([`ConditionalDetrObjectDetectionOutput`]):
            Raw outputs of the model.
        target_sizes (`torch.Tensor` of shape `(batch_size, 2)`):
            Tensor containing the size (h, w) of each image of the batch. For evaluation, this must be the original
            image size (before any data augmentation). For visualization, this should be the image size after data
            augment, but before padding.

Returns:
        `List[Dict]`: A list of dictionaries, each dictionary containing the scores, labels and boxes for an image
        in the batch as predicted by the model.


def __init__(self, config: UnivNetConfig):
    super().__init__()
    self.channels = config.model_in_channels
    self.kernel_size = config.kernel_predictor_conv_size
    self.dropout_prob = config.kernel_predictor_dropout
    self.leaky_relu_slope = config.leaky_relu_slope
    padding = (self.kernel_size - 1) // 2
    self.dropout = nn.Dropout(self.dropout_prob)
    self.conv1 = nn.Conv1d(self.channels, self.channels, self.kernel_size, padding=padding, bias=True)
    self.conv2 = nn.Conv1d(self.channels, self.channels, self.kernel_size, padding=padding, bias=True)

Implementation of the residual block for the kernel predictor network inside each location variable convolution
    block (LVCBlock).

    Parameters:
        config: (`UnivNetConfig`):
            Config for the `UnivNetModel` model.


def call(self, input_ids: TFModelInputType | None=None, attention_mask: np.ndarray | tf.Tensor | None=None, token_type_ids: np.ndarray | tf.Tensor | None=None, position_ids: np.ndarray | tf.Tensor | None=None, head_mask: np.ndarray | tf.Tensor | None=None, inputs_embeds: np.ndarray | tf.Tensor | None=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, next_sentence_label: np.ndarray | tf.Tensor | None=None, training: Optional[bool]=False) -> Union[Tuple, TFNextSentencePredictorOutput]:
    outputs = self.mobilebert(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, training=training)
    pooled_output = outputs[1]
    seq_relationship_scores = self.cls(pooled_output)
    next_sentence_loss = None if next_sentence_label is None else self.hf_compute_loss(labels=next_sentence_label, logits=seq_relationship_scores)
    if not return_dict:
        output = (seq_relationship_scores,) + outputs[2:]
        return (next_sentence_loss,) + output if next_sentence_loss is not None else output
    return TFNextSentencePredictorOutput(loss=next_sentence_loss, logits=seq_relationship_scores, hidden_states=outputs.hidden_states, attentions=outputs.attentions)

Return:

Example:
    ```python
    >>> import tensorflow as tf
    >>> from transformers import AutoTokenizer, TFMobileBertForNextSentencePrediction

    >>> tokenizer = AutoTokenizer.from_pretrained("google/mobilebert-uncased")
    >>> model = TFMobileBertForNextSentencePrediction.from_pretrained("google/mobilebert-uncased")

    >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
    >>> next_sentence = "The sky is blue due to the shorter wavelength of blue light."
    >>> encoding = tokenizer(prompt, next_sentence, return_tensors="tf")

    >>> logits = model(encoding["input_ids"], token_type_ids=encoding["token_type_ids"])[0]
    ```


def run_inside_wrap_function_in_eager_mode(graph_function):

    def wrap_and_execute(self):
        if context.executing_eagerly():
            wrapped = wrap_function.wrap_function(graph_function, [self])
            wrapped()
        else:
            graph_function(self)
    return wrap_and_execute

Decorator to execute the same graph code in eager and graph modes.

    In graph mode, we just execute the graph_function passed as argument. In eager
    mode, we wrap the function using wrap_function and then execute the wrapped
    result.

Args:
        graph_function: python function containing graph code to be wrapped

Returns:
        decorated function


def to_code_v1(entity, recursive=True, arg_values=None, arg_types=None, indentation='  ', experimental_optional_features=None):
    del arg_values
    del arg_types
    del indentation
    return to_code(entity, recursive=recursive, experimental_optional_features=experimental_optional_features)

Returns the source code generated by AutoGraph, as a string.

    Example usage:

    >>> def f(x):
    ...   if x < 0:
    ...     x = -x
    ...   return x
    >>> tf.autograph.to_code(f)
    "...def tf__f(x):..."

    Also see: `tf.autograph.to_graph`.

    Note: If a function has been decorated with `tf.function`, pass its
    underlying Python function, rather than the callable that `tf.function
    creates:

    >>> @tf.function
    ... def f(x):
    ...   if x < 0:
    ...     x = -x
    ...   return x
    >>> tf.autograph.to_code(f.python_function)
    "...def tf__f(x):..."

Args:
        entity: Python callable or class.
        recursive: Whether to recursively convert any functions that the converted
        function may call.
        arg_values: Deprecated.
        arg_types: Deprecated.
        indentation: Deprecated.
        experimental_optional_features: `None`, a tuple of, or a single
        `tf.autograph.experimental.Feature` value.

Returns:
        The converted code as string.


def scan_meta_graph_def(meta_graph_def, op_denylist):
    ops_in_metagraph = set(meta_graph_lib.ops_used_by_graph_def(meta_graph_def.graph_def))
    denylisted_ops = op_denylist & ops_in_metagraph
    if denylisted_ops:
        print('MetaGraph with tag set %s contains the following denylisted ops:' % meta_graph_def.meta_info_def.tags, denylisted_ops)
    else:
        print('MetaGraph with tag set %s does not contain the default denylisted ops:' % meta_graph_def.meta_info_def.tags, op_denylist)

Scans meta_graph_def and reports if there are ops on denylist.

    Print ops if they are on denylist, or print success if no denylisted ops
    found.

Args:
        meta_graph_def: MetaGraphDef protocol buffer.
        op_denylist: set of ops to scan for.


def _run_inline_graph_optimization(func, lower_control_flow, aggressive_inlining):
    graph_def = func.graph.as_graph_def()
    if not lower_control_flow:
        graph_def = disable_lower_using_switch_merge(graph_def)
    for function in graph_def.library.function:
        if 'api_implements' in function.attr:
            del function.attr['api_implements']
    meta_graph = export_meta_graph(graph_def=graph_def, graph=func.graph)
    for name in ['variables', 'model_variables', 'trainable_variables', 'local_variables']:
        raw_list = []
        for raw in meta_graph.collection_def['variables'].bytes_list.value:
            variable = variable_pb2.VariableDef()
            variable.ParseFromString(raw)
            variable.ClearField('initializer_name')
            raw_list.append(variable.SerializeToString())
        meta_graph.collection_def[name].bytes_list.value[:] = raw_list
    fetch_collection = meta_graph_pb2.CollectionDef()
    for array in func.inputs + func.outputs:
        fetch_collection.node_list.value.append(array.name)
    meta_graph.collection_def['train_op'].CopyFrom(fetch_collection)
    config = config_pb2.ConfigProto()
    rewrite_options = config.graph_options.rewrite_options
    rewrite_options.min_graph_nodes = -1
    rewrite_options.optimizers.append('function')
    if aggressive_inlining:
        rewrite_options.function_optimization = rewriter_config_pb2.RewriterConfig.AGGRESSIVE
    return tf_optimizer.OptimizeGraph(config, meta_graph)

Apply function inline optimization to the graph.

    Returns the GraphDef after Grappler's function inlining optimization is
    applied. This optimization does not work on models with control flow.

Args:
        func: ConcreteFunction.
        lower_control_flow: Boolean indicating whether or not to lower control flow
        ops such as If and While. (default True)
        aggressive_inlining: Boolean indicating whether or not to do aggressive
        function inlining (might be unsafe if function has stateful ops not
        properly connected to control outputs).

Returns:
        GraphDef


def np_where(condition, x=None, y=None):
    if x is None and y is None:
        if np.lib.NumpyVersion(np.__version__) >= '2.1.0.rc0':
            return np.atleast_1d(np.asarray(condition)).nonzero()
        return np.where(condition)
    return np.where(condition, x, y)

Return elements chosen from x or y depending on condition.

    When only condition is provided, np.where(condition) is a shorthand for
    np.asarray(condition).nonzero(). See
    https://numpy.org/doc/stable/reference/generated/numpy.where.html. NumPy
    2.1.0rc0 disallows 0D input arrays in nonzero, so np.atleast_1d is used here
    to remain compatible with NumPy 1.x. See
    https://github.com/numpy/numpy/pull/26268.

Args:
        condition: Array_like, bool. Where True, yield x, otherwise yield y.
        x: Array_like. Values from which to choose. x, y and condition need to be
        broadcastable to some shape.
        y: Array_like. Values from which to choose. x, y and condition need to be
        broadcastable to some shape.

Returns:
        An array with elements from x where condition is True, and elements from y
        elsewhere. Or the indices of the elements that are non-zero.


def _track_partially_matched_protocols(self):
    old_protocol_cache = set(self._protocol_cache)
    yield
    self._protocol_cache = old_protocol_cache

Context manager for handling the protocol cache.

    Some protocols have methods that return instances of the protocol, e.g.
    Iterator.next returns Iterator. This will cause an infinite loop, which can
    be avoided by tracking partially matched protocols. To prevent collisions,
    keys are removed from the cache as soon as match is completed.

Yields:
        Into the protocol matching context.


def run(argv=None):
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', dest='input', default='gs://apache-beam-samples/nyc_trip/avro/fhvhv_tripdata_2023-02.avro', help='Input file of NYC FHV data to process. Larger dataset can be found here: gs://apache-beam-samples/nyc_trip/avro/*')
    parser.add_argument('--output', dest='output', help='Output file to write results to.', required=True)
    known_args, pipeline_args = parser.parse_known_args(argv)
    pipeline_options = PipelineOptions(pipeline_args)
    pipeline_options.view_as(SetupOptions).save_main_session = True
    with beam.Pipeline(options=pipeline_options) as p:
        p | ReadFromAvro(known_args.input) | beam.Filter(lambda record: all((record[k] is not None for k in ('hvfhs_license_num', 'trip_miles', 'trip_time', 'base_passenger_fare', 'tips', 'driver_pay'))) and any((record[k] is not None for k in ('request_datetime', 'on_scene_datetime', 'pickup_datetime', 'dropoff_datetime')))) | beam.ParDo(CreateKeyWithServiceAndDay()) | beam.CombinePerKey(CalculatePricePerAttribute()) | beam.Map(flatten_group) | WriteToAvro(known_args.output, SCHEMA, file_name_suffix='.avro')

Runs the New York City trips pipeline.

Args:
        argv: Pipeline options as a list of arguments.


def embedding_lookup(self, features: Any, weights: Optional[Any]=None) -> Any:
    return cpu_embedding_lookup(features, weights, self.embedding_tables, self._feature_config)

Apply standard lookup ops on CPU.

Args:
        features: A nested structure of `tf.Tensor`s, `tf.SparseTensor`s or
        `tf.RaggedTensor`s, with the same structure as `feature_config`. Inputs
        will be downcast to `tf.int32`. Only one type out of `tf.SparseTensor`
        or `tf.RaggedTensor` is supported per call.
        weights: If not `None`, a nested structure of `tf.Tensor`s,
        `tf.SparseTensor`s or `tf.RaggedTensor`s, matching the above, except
        that the tensors should be of float type (and they will be downcast to
        `tf.float32`). For `tf.SparseTensor`s we assume the `indices` are the
        same for the parallel entries from `features` and similarly for
        `tf.RaggedTensor`s we assume the row_splits are the same.

Returns:
        A nested structure of Tensors with the same structure as input features.


def __init__(self, values, reshuffle=True):
    super().__init__()
    if isinstance(values, (str, bytes)):
        raise TypeError('PTransform Create: Refusing to treat string as an iterable. (string=%r)' % values)
    elif isinstance(values, dict):
        values = values.items()
    self.values = tuple(values)
    self.reshuffle = reshuffle
    self._coder = typecoders.registry.get_coder(self.get_output_type())

Initializes a Create transform.

Args:
        values: An object of values for the PCollection


def report_clean(rows):
    print('DCM REPORT CLEAN')
    first = True
    last = False
    for row in rows:
        if row and row[0] == 'Report Fields':
            break
    for row in rows:
        if 'No data returned by the reporting service.' in row:
            break
        if not row or row[0] == 'Grand Total:':
            break
        if first:
            try:
                date_column = row.index('Date')
                row[date_column] = 'Report_Day'
            except ValueError:
                pass
            row = [column_header_sanitize(cell) for cell in row]
        row = ['' if cell.strip() in ('(not set)', '-') else cell for cell in row]
        yield row
        first = False

Helper to fix DCM report issues for BigQuery and ensure schema compliance.

    Memory efficiently cleans each row by fixing:
    * Strips header and footer to preserve only data rows.
    * Changes 'Date' to 'Report_Day' to avoid using reserved name in BigQuery.
    * removes '-' as columns
    * Changes data format to match data studio if datastusio=True.

    Usage example:

    ```
    filename, report = report_file(...)
    rows = report_to_rows(report)
    rows = report_clean(rows)
    ```

Args:
        * rows: (iterator) Rows to clean.

Returns:
        * Iterator of cleaned rows.


def cosh(x):
    if any_symbolic_tensors((x,)):
        return Cosh().symbolic_call(x)
    return backend.numpy.cosh(x)

Hyperbolic cosine, element-wise.

Args:
        x: Input tensor.

Returns:
        Output tensor of same shape as `x`.


def keras_model_summary(name, data, step=None):
    summary_metadata = summary_pb2.SummaryMetadata()
    summary_metadata.plugin_data.plugin_name = 'graph_keras_model'
    summary_metadata.plugin_data.content = b'1'
    try:
        json_string = data.to_json()
    except Exception as exc:
        logging.warning('Model failed to serialize as JSON. Ignoring... %s', exc)
        return False
    with summary_ops_v2.summary_scope(name, 'graph_keras_model', [data, step]) as (tag, _):
        with ops.device('cpu:0'):
            tensor = constant_op.constant(json_string, dtype=dtypes.string)
        return summary_ops_v2.write(tag=tag, tensor=tensor, step=step, metadata=summary_metadata)

Writes a Keras model as JSON to as a Summary.

    Writing the Keras model configuration allows the TensorBoard graph plugin to
    render a conceptual graph, as opposed to graph of ops. In case the model fails
    to serialize as JSON, it ignores and returns False.

Args:
        name: A name for this summary. The summary tag used for TensorBoard will be
        this name prefixed by any active name scopes.
        data: A Keras Model to write.
        step: Explicit `int64`-castable monotonic step value for this summary. If
        omitted, this defaults to `tf.summary.experimental.get_step()`, which must
        not be None.

Returns:
        True on success, or False if no summary was written because no default
        summary writer was available.

Raises:
        ValueError: if a default writer exists, but no step was provided and
        `tf.summary.experimental.get_step()` is None.


def from_json_file(cls, json_file: Union[str, os.PathLike]):
    with open(json_file, encoding='utf-8') as reader:
        text = reader.read()
    image_processor_dict = json.loads(text)
    return cls(**image_processor_dict)

Instantiates a image processor of type [`~image_processing_utils.ImageProcessingMixin`] from the path to a JSON
    file of parameters.

Args:
        json_file (`str` or `os.PathLike`):
            Path to the JSON file containing the parameters.

Returns:
        A image processor of type [`~image_processing_utils.ImageProcessingMixin`]: The image_processor object
        instantiated from that JSON file.


def window_partition(hidden_states, window_size):
    batch_size, height, width, num_channels = hidden_states.shape
    hidden_states = hidden_states.view(batch_size, height // window_size, window_size, width // window_size, window_size, num_channels)
    windows = hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, num_channels)
    return windows

Returns the resized hidden states. The output shape should be `(batch_size * num_windows, window_size, window_size,
    num_channels)`

Args:
        hidden_states (`torch.FloatTensor` of shape `(batch_size, height, width, num_channels)`):
            Input hidden states
        window_size (`int`):
            Window size


def get_size_dict(size: Optional[Union[int, Iterable[int], dict[str, int]]]=None, max_size: Optional[int]=None, height_width_order: bool=True, default_to_square: bool=True, param_name='size') -> dict:
    if not isinstance(size, dict):
        size_dict = convert_to_size_dict(size, max_size, default_to_square, height_width_order)
        logger.info(f'{param_name} should be a dictionary on of the following set of keys: {VALID_SIZE_DICT_KEYS}, got {size}. Converted to {size_dict}.')
    else:
        size_dict = size
    if not is_valid_size_dict(size_dict):
        raise ValueError(f'{param_name} must have one of the following set of keys: {VALID_SIZE_DICT_KEYS}, got {size_dict.keys()}')
    return size_dict

Converts the old size parameter in the config into the new dict expected in the config. This is to ensure backwards
    compatibility with the old image processor configs and removes ambiguity over whether the tuple is in (height,
    width) or (width, height) format.

    - If `size` is tuple, it is converted to `{"height": size[0], "width": size[1]}` or `{"height": size[1], "width":
    size[0]}` if `height_width_order` is `False`.
    - If `size` is an int, and `default_to_square` is `True`, it is converted to `{"height": size, "width": size}`.
    - If `size` is an int and `default_to_square` is False, it is converted to `{"shortest_edge": size}`. If `max_size`
        is set, it is added to the dict as `{"longest_edge": max_size}`.

Args:
        size (`Union[int, Iterable[int], Dict[str, int]]`, *optional*):
            The `size` parameter to be cast into a size dictionary.
        max_size (`Optional[int]`, *optional*):
            The `max_size` parameter to be cast into a size dictionary.
        height_width_order (`bool`, *optional*, defaults to `True`):
            If `size` is a tuple, whether it's in (height, width) or (width, height) order.
        default_to_square (`bool`, *optional*, defaults to `True`):
            If `size` is an int, whether to default to a square image or not.


def get_visible_devices(device_type=None):
    return context.context().get_visible_devices(device_type)

Get the list of visible physical devices.

    Returns the list of `PhysicalDevice`s currently marked as visible to the
    runtime. A visible device will have at least one `LogicalDevice` associated
    with it once the runtime is initialized.

    The following example verifies all visible GPUs have been disabled:

    >>> physical_devices = tf.config.list_physical_devices('GPU')
    >>> try:
    ...   # Disable all GPUS
    ...   tf.config.set_visible_devices([], 'GPU')
    ...   visible_devices = tf.config.get_visible_devices()
    ...   for device in visible_devices:
    ...     assert device.device_type != 'GPU'
    ... except:
    ...   # Invalid device or cannot modify virtual devices once initialized.
    ...   pass

Args:
        device_type: (optional string) Only include devices matching this device
        type. For example "CPU" or "GPU".

Returns:
        List of visible `PhysicalDevice`s


def _aggregate_grads(gradients):
    assert gradients, 'No gradients to aggregate'
    if len(gradients) == 1:
        return gradients[0]
    if all((isinstance(g, tensor_lib.Tensor) for g in gradients)):
        return gen_math_ops.add_n(gradients)
    else:
        assert all((isinstance(g, (tensor_lib.Tensor, indexed_slices.IndexedSlices)) for g in gradients))
        return backprop_util.AggregateIndexedSlicesGradients(gradients)

Aggregate gradients from multiple sources.

Args:
        gradients: A list of 'Tensor' or 'IndexedSlices' gradients.

Returns:
        If 'gradients' only has 'Tensor', returns an aggregated 'Tensor'.
        Otherwise returns an aggregated 'IndexedSlices'.


def _least_upper_bound(*nodes):
    N = set(nodes)
    UB = LATTICE_UPPER_BOUNDS
    try:
        bounds = [UB[n] for n in N]
    except KeyError:
        dtype = next((n for n in N if n not in UB))
        raise ValueError(f'dtype={dtype!r} is not a valid dtype for Keras type promotion.')
    CUB = set.intersection(*bounds)
    LUB = CUB & N or {c for c in CUB if CUB.issubset(UB[c])}
    if len(LUB) == 1:
        return LUB.pop()
    elif len(LUB) == 0:
        msg = f'Input dtypes {tuple((str(n) for n in nodes))} have no available implicit dtype promotion path. Try explicitly casting inputs to the desired output type.'
        raise ValueError(msg)
    else:
        raise ValueError(f"Internal Type Promotion error: {nodes} do not have a unique least upper bound on the specified lattice; options are {LUB}. This is an unexpected error in Keras's internal logic; please report it to the maintainers.")

Compute the least upper bound of a set of nodes.

Args:
        nodes: sequence of entries from dtypes + weak_types

Returns:
        The type representing the least upper bound of the input nodes on the
        promotion lattice.


def __init__(self, dump_root, partition_graphs=None, validate=True):
    if not gfile.IsDirectory(dump_root):
        raise IOError('Dump root directory %s does not exist' % dump_root)
    self._core_metadata = []
    self._dump_root = dump_root
    self._load_core_metadata()
    self._load_fetches_info()
    self._load_feeds_info()
    self._load_all_device_dumps(partition_graphs, validate)
    self._python_graph = None

`DebugDumpDir` constructor.

Args:
        dump_root: (`str`) path to the dump root directory.
        partition_graphs: A repeated field of GraphDefs representing the
        partition graphs executed by the TensorFlow runtime.
        validate: (`bool`) whether the dump files are to be validated against the
        partition graphs.

Raises:
        IOError: If dump_root does not exist as a directory.
        ValueError: If more than one core metadata file is found under the dump
        root directory.


def get_float(self, min_float=_MIN_FLOAT, max_float=_MAX_FLOAT):
    return self.fdp.ConsumeFloatInRange(min_float, max_float)

Consume a float with given constraints.

Args:
        min_float: Minimum allowed float.
        max_float: Maximum allowed float.

Returns:
        Consumed float based on input bytes and constraints.


def _ContainsAll(self, verb, expected):
    actual_list = list(self._actual)
    missing = _DuplicateCounter()
    actual_not_in_order = set()
    ordered = True
    for i in expected:
        try:
            index = actual_list.index(i)
            for _ in six.moves.xrange(index):
                actual_element = actual_list.pop(0)
                if _IsHashable(actual_element) and isinstance(actual_not_in_order, collections_abc.Set):
                    actual_not_in_order.add(actual_element)
                else:
                    if isinstance(actual_not_in_order, collections_abc.Set):
                        actual_not_in_order = list(actual_not_in_order)
                    if actual_element not in actual_not_in_order:
                        actual_not_in_order.append(actual_element)
            actual_list.pop(0)
        except ValueError:
            if not _IsHashable(i) and isinstance(actual_not_in_order, collections_abc.Set):
                actual_not_in_order = list(actual_not_in_order)
            if i in actual_not_in_order:
                actual_not_in_order.remove(i)
                ordered = False
            else:
                missing.Increment(i)
    if missing:
        self._FailWithBadResults(verb, expected, 'is missing', missing)
    if ordered:
        return _InOrder()
    else:
        return _NotInOrder(self._actual, 'contains all elements in order', expected)

Determines if the subject contains all the expected elements.

    Helper function for ContainsAllIn() and ContainsAllOf().

Args:
        verb: string describing how the expected elements should be contained.
        expected: iterable of objects that should be contained in the subject.

Returns:
        If the subject does contain all the expected elements, returns an
        _Ordered predicate on which .InOrder() can be subsequently called.

Raises:
        TruthAssertionError: the subject is missing any of the expected elements.


def convert_to_rgb(self, video: 'torch.Tensor') -> VideoInput:
    video = F.grayscale_to_rgb(video)
    if video.shape[-3] == 3 or not (video[..., 3, :, :] < 255).any():
        return video
    alpha = video[..., 3, :, :] / 255.0
    video = (1 - alpha[..., None, :, :]) * 255 + alpha[..., None, :, :] * video[..., :3, :, :]
    return video

Converts a video to RGB format.

Args:
        video (`"torch.Tensor"`):
            The video to convert.

Returns:
        `torch.Tensor`: The converted video.


def get_profiles(self, cmd):
    if cmd not in self._views:
        raise ValueError('No autoprofiler for command: {}, was run'.format(cmd))
    return self._views[cmd]

Returns profiling results for each step at which `cmd` was run.

Args:
        cmd: string, profiling command used in an `add_auto_profiling` call.

Returns:
        dict[int: (MultiGraphNodeProto | GraphNodeProto)]. Keys are steps at which
        the profiling command was run. Values are the outputs of profiling.
        For "code" and "op" commands this will be a `MultiGraphNodeProto`, for
        "scope" and "graph" commands this will be a `GraphNodeProto.

Raises:
        ValueError: if `cmd` was never run (either because no session.run call was
        made or because there was no `add_auto_profiling` call with the specified
        `cmd`.


def _create_partition_config(option: t.Tuple, config: Config) -> Config:
    copy = cp.deepcopy(config.selection)
    out = cp.deepcopy(config)
    for idx, key in enumerate(config.partition_keys):
        copy[key] = [option[idx]]
    if 'hdate' in copy:
        copy['hdate'] = [generate_hdate(copy['date'][0], v) for v in copy['hdate']]
    out.selection = copy
    return out

Create a config for a single partition option.

    Output a config dictionary, overriding the range of values for
    each key with the partition instance in 'selection'.
    Continuing the example from prepare_partitions, the selection section
    would be:
        { 'foo': ..., 'year': ['2020'], 'month': ['01'], ... }
        { 'foo': ..., 'year': ['2020'], 'month': ['02'], ... }
        { 'foo': ..., 'year': ['2020'], 'month': ['03'], ... }

Args:
        option: A single item in the range of partition_keys.
        config: The download config, including the parameters and selection sections.

Returns:
        A configuration with that selects a single download partition.


def __init__(self, num_rows, multiplier, is_non_singular=None, is_self_adjoint=None, is_positive_definite=None, is_square=True, assert_proper_shapes=False, name='LinearOperatorScaledIdentity'):
    parameters = dict(num_rows=num_rows, multiplier=multiplier, is_non_singular=is_non_singular, is_self_adjoint=is_self_adjoint, is_positive_definite=is_positive_definite, is_square=is_square, assert_proper_shapes=assert_proper_shapes, name=name)
    self._assert_proper_shapes = assert_proper_shapes
    with ops.name_scope(name, values=[multiplier, num_rows]):
        self._multiplier = linear_operator_util.convert_nonref_to_tensor(multiplier, name='multiplier')
        if not self._multiplier.dtype.is_complex:
            if is_self_adjoint is False:
                raise ValueError('A real diagonal operator is always self adjoint.')
            else:
                is_self_adjoint = True
        if not is_square:
            raise ValueError('A ScaledIdentity operator is always square.')
        linear_operator_util.assert_not_ref_type(num_rows, 'num_rows')
        super(LinearOperatorScaledIdentity, self).__init__(dtype=self._multiplier.dtype.base_dtype, is_non_singular=is_non_singular, is_self_adjoint=is_self_adjoint, is_positive_definite=is_positive_definite, is_square=is_square, parameters=parameters, name=name)
        self._num_rows = linear_operator_util.shape_tensor(num_rows, name='num_rows')
        self._num_rows_static = tensor_util.constant_value(self._num_rows)
        self._check_num_rows_possibly_add_asserts()
        self._num_rows_cast_to_dtype = math_ops.cast(self._num_rows, self.dtype)
        self._num_rows_cast_to_real_dtype = math_ops.cast(self._num_rows, self.dtype.real_dtype)

Initialize a `LinearOperatorScaledIdentity`.

    The `LinearOperatorScaledIdentity` is initialized with `num_rows`, which
    determines the size of each identity matrix, and a `multiplier`,
    which defines `dtype`, batch shape, and scale of each matrix.

    This operator is able to broadcast the leading (batch) dimensions.

Args:
        num_rows:  Scalar non-negative integer `Tensor`.  Number of rows in the
        corresponding identity matrix.
        multiplier:  `Tensor` of shape `[B1,...,Bb]`, or `[]` (a scalar).
        is_non_singular:  Expect that this operator is non-singular.
        is_self_adjoint:  Expect that this operator is equal to its hermitian
        transpose.
        is_positive_definite:  Expect that this operator is positive definite,
        meaning the quadratic form `x^H A x` has positive real part for all
        nonzero `x`.  Note that we do not require the operator to be
        self-adjoint to be positive-definite.  See:
        https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
        is_square:  Expect that this operator acts like square [batch] matrices.
        assert_proper_shapes:  Python `bool`.  If `False`, only perform static
        checks that initialization and method arguments have proper shape.
        If `True`, and static checks are inconclusive, add asserts to the graph.
        name: A name for this `LinearOperator`

Raises:
        ValueError:  If `num_rows` is determined statically to be non-scalar, or
        negative.


def _compute_direction_numbers(dim):
    m = np.empty((dim, 32), dtype=np.int32)
    m[0, :] = np.ones(32, dtype=np.int32)
    for k in range(dim - 1):
        a_k = _PRIMITIVE_POLYNOMIAL_COEFFICIENTS[k]
        deg = np.int32(np.floor(np.log2(a_k)))
        m[k + 1, :deg] = _INITIAL_DIRECTION_NUMBERS[:deg, k]
        for j in range(deg, 32):
            m[k + 1, j] = m[k + 1, j - deg]
            for i in range(deg):
                if a_k >> i & 1:
                    m[k + 1, j] = np.bitwise_xor(m[k + 1, j], m[k + 1, j - deg + i] << deg - i)
    return m

Returns array of direction numbers for dimension dim.

    These are the m_kj values in the Joe & Kuo notes[1], not the v_kj values. So
    these refer to the 'abuse of notation' mentioned in the notes -- it is a
    matrix of integers, not floats. The variable names below are intended to match
    the notation in the notes as closely as possible.

Args:
        dim: int, dimension.

Returns:
        `numpy.array` of direction numbers with `shape` [dim, 32].


def _get_help_for_command_prefix(self, cmd_prefix):
    lines = []
    resolved_prefix = self._resolve_prefix(cmd_prefix)
    if not resolved_prefix:
        lines.append('Invalid command prefix: "%s"' % cmd_prefix)
        return lines
    lines.append(resolved_prefix)
    if resolved_prefix in self._prefix_to_aliases:
        lines.append(HELP_INDENT + 'Aliases: ' + ', '.join(self._prefix_to_aliases[resolved_prefix]))
    lines.append('')
    help_lines = self._prefix_to_help[resolved_prefix].split('\n')
    for line in help_lines:
        lines.append(HELP_INDENT + line)
    return lines

Compile the help information for a given command prefix.

Args:
        cmd_prefix: Command prefix, as the prefix itself or one of its aliases.

Returns:
        A list of str as the help information for cmd_prefix. If the cmd_prefix
        does not exist, the returned list of str will indicate that.


def __init__(self, initial_learning_rate, first_decay_steps, t_mul=2.0, m_mul=1.0, alpha=0.0, name=None):
    super(CosineDecayRestarts, self).__init__()
    self.initial_learning_rate = initial_learning_rate
    self.first_decay_steps = first_decay_steps
    self._t_mul = t_mul
    self._m_mul = m_mul
    self.alpha = alpha
    self.name = name

Applies cosine decay with restarts to the learning rate.

Args:
        initial_learning_rate: A scalar `float32` or `float64` Tensor or a Python
        number. The initial learning rate.
        first_decay_steps: A scalar `int32` or `int64` `Tensor` or a Python
        number. Number of steps to decay over.
        t_mul: A scalar `float32` or `float64` `Tensor` or a Python number.
        Used to derive the number of iterations in the i-th period
        m_mul: A scalar `float32` or `float64` `Tensor` or a Python number.
        Used to derive the initial learning rate of the i-th period:
        alpha: A scalar `float32` or `float64` Tensor or a Python number.
        Minimum learning rate value as a fraction of the initial_learning_rate.
        name: String. Optional name of the operation.  Defaults to 'SGDRDecay'.


def forward(self, hidden_states, output_router_logits):
    forwarded_states, router_tuple = self.mlp(hidden_states)
    forwarded_states += torch.tanh(self.soft_bypass_mlp(hidden_states))
    output = hidden_states + self.norm(forwarded_states)
    if output_router_logits and router_tuple is not None:
        return (output, router_tuple)
    else:
        return output

Args:
    hidden_states (`torch.Tensor`) :
        [num_groups, tokens_per_group, hidden_dim] inputs to send to experts.
    output_router_logits (`bool`) :
        output experts router output.

Returns:
    torch.Tensor[num_groups, tokens_per_group, hidden_dim]


def _cast_indexed_slice_indices(a, b):
    if isinstance(a, indexed_slices.IndexedSlices) and isinstance(b, indexed_slices.IndexedSlices) and (a.indices.dtype != b.indices.dtype):
        a._indices = math_ops.cast(a.indices, dtypes.int64)
        b._indices = math_ops.cast(b.indices, dtypes.int64)

Cast IndexedSlice.indices from int32 to int64 where necessary.

    If `a` and `b` are both IndexedSlices, and their indices have different
    dtypes, then cast both their dtypes to `int64` (modifies `a` and `b`
    in-place).  Otherwise, does nothing.

Args:
        a: A value, which may be an IndexedSlices.
        b: A value, which may be an IndexedSlices.


def center_crop(self, image: np.ndarray, size: dict[str, int], data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None, **kwargs) -> np.ndarray:
    size = get_size_dict(size)
    if 'height' not in size or 'width' not in size:
        raise ValueError(f"The size dictionary must have keys 'height' and 'width'. Got {size.keys()}")
    return center_crop(image, size=(size['height'], size['width']), data_format=data_format, input_data_format=input_data_format, **kwargs)

Center crop an image to `(size["height"], size["width"])`. If the input size is smaller than `crop_size` along
    any edge, the image is padded with 0's and then center cropped.

Args:
        image (`np.ndarray`):
            Image to center crop.
        size (`Dict[str, int]`):
            Size of the output image.
        data_format (`str` or `ChannelDimension`, *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.
        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.


def forward(self, input_features: Optional[torch.FloatTensor]=None, is_longer: Optional[torch.BoolTensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple, BaseModelOutputWithPooling]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return self.audio_encoder(input_features=input_features, is_longer=is_longer, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)

input_features (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
    Input audio features. This should be returned by the [`ClapFeatureExtractor`] class that you can also
    retrieve from [`AutoFeatureExtractor`]. See [`ClapFeatureExtractor.__call__`] for details.
    is_longer (`torch.FloatTensor`, of shape `(batch_size, 1)`, *optional*):
    Whether the audio clip is longer than `max_length`. If `True`, a feature fusion will be enabled to enhance
    the features.

Example:
    ```python
    >>> from datasets import load_dataset
    >>> from transformers import AutoProcessor, ClapAudioModel

    >>> dataset = load_dataset("hf-internal-testing/ashraq-esc50-1-dog-example")
    >>> audio_sample = dataset["train"]["audio"][0]["array"]

    >>> model = ClapAudioModel.from_pretrained("laion/clap-htsat-fused")
    >>> processor = AutoProcessor.from_pretrained("laion/clap-htsat-fused")

    >>> inputs = processor(audios=audio_sample, return_tensors="pt")

    >>> outputs = model(**inputs)
    >>> last_hidden_state = outputs.last_hidden_state
    ```


def __init__(self, feature_config: Union[tpu_embedding_v2_utils.FeatureConfig, Iterable], optimizer: Optional[tpu_embedding_v2_utils._Optimizer], pipeline_execution_with_tensor_core: bool=False):
    self._strategy = distribute_lib.get_strategy()
    self._using_tpu = isinstance(self._strategy, (tpu_strategy.TPUStrategy, tpu_strategy.TPUStrategyV2))
    self._pipeline_execution_with_tensor_core = pipeline_execution_with_tensor_core
    self._feature_config = feature_config
    self._output_shapes = []
    for feature in nest.flatten(feature_config):
        self._output_shapes.append(feature.output_shape)
    device_assignment = getattr(self._strategy.extended, '_device_assignment', None)
    self._num_cores_per_replica = device_assignment.num_cores_per_replica if device_assignment else None
    self._table_config = []
    for feature in nest.flatten(feature_config):
        if feature.table not in self._table_config:
            self._table_config.append(feature.table)
    table_names = []
    for i, table in enumerate(self._table_config):
        if table.optimizer is None:
            table.optimizer = optimizer
        if (table.optimizer is not None or self._using_tpu) and (not isinstance(table.optimizer, tpu_embedding_v2_utils._Optimizer)):
            raise ValueError('{} is an unsupported optimizer class. Please pass an instance of one of the optimizer classes under tf.tpu.experimental.embedding.'.format(type(table.optimizer)))
        if table.name is None:
            table.name = 'table_{}'.format(i)
        if table.name in table_names:
            raise ValueError(f'Tables must have a unique name. Multiple tables with name {table.name} found.')
        table_names.append(table.name)
    if self._using_tpu:
        self._dynamic_learning_rates = []
        for table in self._table_config:
            if callable(table.optimizer.learning_rate) and table.optimizer.learning_rate not in self._dynamic_learning_rates:
                self._dynamic_learning_rates.append(table.optimizer.learning_rate)
        self._hosts = tpu_embedding_v2_utils.get_list_of_hosts(self._strategy)
    self._built = False
    self._verify_output_shapes_on_enqueue = True

Creates the TPUEmbedding mid level API object.

    ```python
    strategy = tf.distribute.TPUStrategy(...)
    with strategy.scope():
        embedding = tf.tpu.experimental.embedding.TPUEmbedding(
        feature_config=tf.tpu.experimental.embedding.FeatureConfig(
            table=tf.tpu.experimental.embedding.TableConfig(
                dim=...,
                vocabulary_size=...)))
    ```

Args:
        feature_config: A nested structure of
        `tf.tpu.experimental.embedding.FeatureConfig` configs.
        optimizer: An instance of one of `tf.tpu.experimental.embedding.SGD`,
        `tf.tpu.experimental.embedding.Adagrad` or
        `tf.tpu.experimental.embedding.Adam`. When not created under
        TPUStrategy may be set to None to avoid the creation of the optimizer
        slot variables, useful for optimizing memory consumption when exporting
        the model for serving where slot variables aren't needed.
        pipeline_execution_with_tensor_core: If True, the TPU embedding
        computations will overlap with the TensorCore computations (and hence
        will be one step old). Set to True for improved performance.

Raises:
        ValueError: If optimizer is not one of tf.tpu.experimental.embedding.(SGD,
        Adam or Adagrad) or None when created under a TPUStrategy.


def delete_directory(directory):
    filesystems.FileSystems.delete([directory])

Delete a directory in a filesystem.

Args:
        directory: Full path to a directory supported by Beam filesystems (e.g.
        "gs://mybucket/mydir/", "s3://...", ...)


def forward(self, inputs_embeds, attention_mask: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None) -> BaseModelOutput:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    encoder_states = () if output_hidden_states else None
    all_attentions = () if output_attentions else None
    hidden_states = inputs_embeds
    for encoder_layer in self.layers:
        if output_hidden_states:
            encoder_states = encoder_states + (hidden_states,)
        layer_outputs = encoder_layer(hidden_states, attention_mask, output_attentions=output_attentions)
        hidden_states = layer_outputs[0]
        if output_attentions:
            all_attentions = all_attentions + (layer_outputs[1],)
    if output_hidden_states:
        encoder_states = encoder_states + (hidden_states,)
    return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions)

Args:
    inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
        Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
        This is useful if you want more control over how to convert `input_ids` indices into associated vectors
        than the model's internal embedding lookup matrix.
    attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
        Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

        - 1 for tokens that are **not masked**,
        - 0 for tokens that are **masked**.

        [What are attention masks?](../glossary#attention-mask)
    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.
    output_hidden_states (`bool`, *optional*):
        Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
        for more detail.
    return_dict (`bool`, *optional*):
        Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.


def _find_op(graph: ops.Graph, op_name: Optional[str]) -> Optional[ops.Operation]:
    if not op_name:
        return None
    init_op = graph.get_operation_by_name(op_name)
    logging.debug('Op found in the graph: %s', op_name)
    return init_op

Finds the operation with `op_name`.

Args:
        graph: The graph to find from.
        op_name: Name of the node.

Returns:
        The operation that corresponds to `op_name`. Returns None iff op_name is an
        empty string or None.

Raises:
        ValueError: `op_name` is malformed.


def squeeze(x, axis=None):
    if any_symbolic_tensors((x,)):
        return Squeeze(axis=axis).symbolic_call(x)
    return backend.numpy.squeeze(x, axis=axis)

Remove axes of length one from `x`.

Args:
        x: Input tensor.
        axis: Select a subset of the entries of length one in the shape.

Returns:
        The input tensor with all or a subset of the dimensions of
        length 1 removed.


def from_json(cls, json_value: Dict[str, Any], *, allow_partial: bool=False, root_path: Optional[utils.KeyPath]=None) -> 'DNA':
    cloneable_metadata_keys = json_value.pop('_cloneable_metadata_keys', None)
    if json_value.get('format', None) == 'compact':
        with symbolic.enable_type_check(False):
            dna = DNA.parse(symbolic.from_json(json_value.get('value')))
            if 'metadata' in json_value:
                dna.rebind(metadata=symbolic.from_json(json_value.get('metadata')), raise_on_no_change=False, skip_notification=True)
    else:
        dna = super(DNA, cls).from_json(json_value, allow_partial=allow_partial, root_path=root_path)
        assert isinstance(dna, DNA)
    if cloneable_metadata_keys:
        dna._cloneable_metadata_keys = set(cloneable_metadata_keys)
    return dna

Class method that load a DNA from a JSON value.

Args:
        json_value: Input JSON value, only JSON dict is acceptable.
        allow_partial: Whether to allow elements of the list to be partial.
        root_path: KeyPath of loaded object in its object tree.

Returns:
        A DNA object.


def _save_model(self):
    if not file_utils.exists(self.backup_dir):
        file_utils.makedirs(self.backup_dir)
    if self.double_checkpoint and file_utils.exists(self._weights_path):
        file_utils.copy(self._weights_path, self._prev_weights_path)
    if self.double_checkpoint and file_utils.exists(self._training_metadata_path):
        file_utils.copy(self._training_metadata_path, self._prev_training_metadata_path)
    self.model.save_weights(filepath=self._weights_path, overwrite=True)
    with file_utils.File(self._training_metadata_path, 'w') as f:
        training_metadata = {'epoch': self._current_epoch, 'batch': self._last_batch_seen}
        f.write(json.dumps(training_metadata))

Saves the model.

Args:
        epoch: the epoch this iteration is in.
        batch: the batch this iteration is in. `None` if the `save_freq`
            is set to `"epoch"`.
        logs: the `logs` dict passed in to `on_batch_end` or `on_epoch_end`.


def __call__(self, x: core.Tensor) -> Mapping[str, core.Tensor]:
    out = nn_ops.conv2d(x, self.filters, strides=[1, 1, 2, 1], dilations=[1, 1, 1, 1], padding='SAME', data_format='NHWC')
    return {'output': out}

Performs a 2D convolution operation.

Args:
        x: Input tensor to perform convolution on.

Returns:
        A map of: output key -> output result.


def __init__(self, packages: Optional[List[FhirPackage[_StructDefT, _SearchParameterT, _CodeSystemT, _ValueSetT]]]=None) -> None:
    self.packages = packages or []

Manages access to a collection of FhirPackage instances.

    Allows users to add packages to the package manager and then search all of
    them for a particular resource.

Attributes:
        packages: The packages added to the package manager.


def forward(self, outputs, targets):
    batch_size, num_queries = outputs['logits'].shape[:2]
    out_prob = outputs['logits'].flatten(0, 1).softmax(-1)
    out_bbox = outputs['pred_boxes'].flatten(0, 1)
    target_ids = torch.cat([v['class_labels'] for v in targets])
    target_bbox = torch.cat([v['boxes'] for v in targets])
    class_cost = -out_prob[:, target_ids]
    bbox_cost = torch.cdist(out_bbox, target_bbox, p=1)
    giou_cost = -generalized_box_iou(center_to_corners_format(out_bbox), center_to_corners_format(target_bbox))
    cost_matrix = self.bbox_cost * bbox_cost + self.class_cost * class_cost + self.giou_cost * giou_cost
    cost_matrix = cost_matrix.view(batch_size, num_queries, -1).cpu()
    sizes = [len(v['boxes']) for v in targets]
    indices = [linear_sum_assignment(c[i]) for i, c in enumerate(cost_matrix.split(sizes, -1))]
    return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]

Args:
    outputs (`dict`):
        A dictionary that contains at least these entries:
        * "logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits
        * "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates.
    targets (`List[dict]`):
        A list of targets (len(targets) = batch_size), where each target is a dict containing:
        * "class_labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of
        ground-truth
        objects in the target) containing the class labels
        * "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates.

Returns:
    `List[Tuple]`: A list of size `batch_size`, containing tuples of (index_i, index_j) where:
    - index_i is the indices of the selected predictions (in order)
    - index_j is the indices of the corresponding selected targets (in order)
    For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes)


def pred(scores: jax.Array, rows: jax.Array, cols: jax.Array, N: int) -> jax.Array:
    r: jax.Array = 2 * jax.ops.segment_sum(scores.take(cols), rows, N) - scores.sum()
    return r > 0

Predicts the target output from the learned scores and input entries.

Args:
        scores (jax.Array): Contribution scores of features.
        rows (jax.Array): Row indices of True values in the input.
        cols (jax.Array): Column indices of True values in the input.
        N (int): The number of input entries.

Returns:
        res (jax.Array): A prediction of the target.


def flip_back(output_flipped, flip_pairs, target_type='gaussian-heatmap'):
    if target_type not in ['gaussian-heatmap', 'combined-target']:
        raise ValueError('target_type should be gaussian-heatmap or combined-target')
    if output_flipped.ndim != 4:
        raise ValueError('output_flipped should be [batch_size, num_keypoints, height, width]')
    batch_size, num_keypoints, height, width = output_flipped.shape
    channels = 1
    if target_type == 'combined-target':
        channels = 3
        output_flipped[:, 1::3, ...] = -output_flipped[:, 1::3, ...]
    output_flipped = output_flipped.reshape(batch_size, -1, channels, height, width)
    output_flipped_back = output_flipped.clone()
    for left, right in flip_pairs.tolist():
        output_flipped_back[:, left, ...] = output_flipped[:, right, ...]
        output_flipped_back[:, right, ...] = output_flipped[:, left, ...]
    output_flipped_back = output_flipped_back.reshape((batch_size, num_keypoints, height, width))
    output_flipped_back = output_flipped_back.flip(-1)
    return output_flipped_back

Flip the flipped heatmaps back to the original form.

Args:
        output_flipped (`torch.tensor` of shape `(batch_size, num_keypoints, height, width)`):
            The output heatmaps obtained from the flipped images.
        flip_pairs (`torch.Tensor` of shape `(num_keypoints, 2)`):
            Pairs of keypoints which are mirrored (for example, left ear -- right ear).
        target_type (`str`, *optional*, defaults to `"gaussian-heatmap"`):
            Target type to use. Can be gaussian-heatmap or combined-target.
            gaussian-heatmap: Classification target with gaussian distribution.
            combined-target: The combination of classification target (response map) and regression target (offset map).
            Paper ref: Huang et al. The Devil is in the Details: Delving into Unbiased Data Processing for Human Pose Estimation (CVPR 2020).

Returns:
        torch.Tensor: heatmaps that flipped back to the original image


def __init__(self, pretrained_cfg: Dict[str, Any], architecture: Optional[str]=None, **kwargs):
    requires_backends(self, 'timm')
    super().__init__(architecture=architecture)
    self.data_config = timm.data.resolve_data_config(pretrained_cfg, model=None, verbose=False)
    self.val_transforms = timm.data.create_transform(**self.data_config, is_training=False)
    self.train_transforms = timm.data.create_transform(**self.data_config, is_training=True)
    self._not_supports_tensor_input = any((transform.__class__.__name__ == 'ToTensor' for transform in self.val_transforms.transforms))

Wrapper class for timm models to be used within transformers.

Args:
        pretrained_cfg (`Dict[str, Any]`):
            The configuration of the pretrained model used to resolve evaluation and
            training transforms.
        architecture (`Optional[str]`, *optional*):
            Name of the architecture of the model.


def __init__(self, model_name: str, columns: list[str], api_key: Optional[str]=None, organization: Optional[str]=None, dimensions: Optional[int]=None, user: Optional[str]=None, max_batch_size: Optional[int]=None, **kwargs):
    self.model_name = model_name
    self.api_key = api_key
    self.organization = organization
    self.dimensions = dimensions
    self.user = user
    self.max_batch_size = max_batch_size
    super().__init__(columns=columns, **kwargs)

Embedding Config for OpenAI Text Embedding models.
    Text Embeddings are generated for a batch of text using the OpenAI API.

Args:
        model_name: Name of the OpenAI embedding model
        columns: The columns where the embeddings will be stored in the output
        api_key: OpenAI API key
        organization: OpenAI organization ID
        dimensions: Specific embedding dimensions to use (if model supports it)
        user: End-user identifier for tracking and rate limit calculations
        max_batch_size: Maximum batch size for requests to OpenAI API


def _safe_exec_func(self, func, *args):
    try:
        return func(*args)
    except signals.TestAbortAll:
        raise
    except Exception:
        logging.exception('Exception happened when executing %s in %s.', func.__name__, self.TAG)

Executes a function with exception safeguard.

    This will let signals.TestAbortAll through so abort_all works in all
    procedure functions.

Args:
        func: Function to be executed.
        args: Arguments to be passed to the function.

Returns:
        Whatever the function returns.


def run_model_pipeline_tests(self, task, repo_name, model_architecture, tokenizer_names, image_processor_names, feature_extractor_names, processor_names, commit, torch_dtype='float32'):
    pipeline_test_class_name = pipeline_test_mapping[task]['test'].__name__
    tokenizer_names = tokenizer_names or [None]
    image_processor_names = image_processor_names or [None]
    feature_extractor_names = feature_extractor_names or [None]
    processor_names = processor_names or [None]
    test_cases = [{'tokenizer_name': tokenizer_name, 'image_processor_name': image_processor_name, 'feature_extractor_name': feature_extractor_name, 'processor_name': processor_name} for tokenizer_name in tokenizer_names for image_processor_name in image_processor_names for feature_extractor_name in feature_extractor_names for processor_name in processor_names]
    for test_case in test_cases:
        tokenizer_name = test_case['tokenizer_name']
        image_processor_name = test_case['image_processor_name']
        feature_extractor_name = test_case['feature_extractor_name']
        processor_name = test_case['processor_name']
        do_skip_test_case = self.is_pipeline_test_to_skip(pipeline_test_class_name, model_architecture.config_class, model_architecture, tokenizer_name, image_processor_name, feature_extractor_name, processor_name)
        if do_skip_test_case:
            logger.warning(f'{self.__class__.__name__}::test_pipeline_{task.replace('-', '_')}_{torch_dtype} is skipped: test is currently known to fail for: model `{model_architecture.__name__}` | tokenizer `{tokenizer_name}` | image processor `{image_processor_name}` | feature extractor {feature_extractor_name}.')
            continue
        self.run_pipeline_test(task, repo_name, model_architecture, tokenizer_name=tokenizer_name, image_processor_name=image_processor_name, feature_extractor_name=feature_extractor_name, processor_name=processor_name, commit=commit, torch_dtype=torch_dtype)

Run pipeline tests for a specific `task` with the give model class and tokenizer/processor class names

Args:
        task (`str`):
            A task name. This should be a key in the mapping `pipeline_test_mapping`.
        repo_name (`str`):
            A model repository id on the Hub.
        model_architecture (`type`):
            A subclass of `PretrainedModel` or `PretrainedModel`.
        tokenizer_names (`List[str]`):
            A list of names of a subclasses of `PreTrainedTokenizerFast` or `PreTrainedTokenizer`.
        image_processor_names (`List[str]`):
            A list of names of subclasses of `BaseImageProcessor`.
        feature_extractor_names (`List[str]`):
            A list of names of subclasses of `FeatureExtractionMixin`.
        processor_names (`List[str]`):
            A list of names of subclasses of `ProcessorMixin`.
        commit (`str`):
            The commit hash of the model repository on the Hub.
        torch_dtype (`str`, `optional`, defaults to `'float32'`):
            The torch dtype to use for the model. Can be used for FP16/other precision inference.


Returns the object holds the test environment information.

    Tests should modify this in the main process if needed, and it will be passed
    to the worker processes each time a test case is run.

Returns:
        a TestEnvironment object.


def __init__(self, counters=None, distributions=None, gauges=None, string_sets=None, bounded_tries=None):
    self.counters = counters or {}
    self.distributions = distributions or {}
    self.gauges = gauges or {}
    self.string_sets = string_sets or {}
    self.bounded_tries = bounded_tries or {}

Create a MetricUpdates object.

Args:
        counters: Dictionary of MetricKey:MetricUpdate updates.
        distributions: Dictionary of MetricKey:MetricUpdate objects.
        gauges: Dictionary of MetricKey:MetricUpdate objects.
        string_sets: Dictionary of MetricKey:MetricUpdate objects.
        bounded_tries: Dictionary of MetricKey:MetricUpdate objects.


def __init__(self, dataset=None, components=None, element_spec=None):
    super(OwnedIterator, self).__init__()
    if dataset is None:
        if components is None or element_spec is None:
            raise ValueError('When `dataset` is not provided, both `components` and `element_spec` must be specified.')
        self._element_spec = element_spec
        self._flat_output_types = structure.get_flat_tensor_types(self._element_spec)
        self._flat_output_shapes = structure.get_flat_tensor_shapes(self._element_spec)
        self._components = components
        self._iterator_resource, = components
    else:
        if components is not None or element_spec is not None:
            raise ValueError('When `dataset` is provided, `element_spec` and `components` must not be specified.')
        self._create_iterator(dataset)
    self._get_next_call_count = 0

Creates a new iterator from the given dataset.

    If `dataset` is not specified, the iterator will be created from the given
    tensor components and element structure. In particular, the alternative for
    constructing the iterator is used when the iterator is reconstructed from
    it `CompositeTensor` representation.

Args:
        dataset: A `tf.data.Dataset` object.
        components: Tensor components to construct the iterator from.
        element_spec: A (nested) structure of `TypeSpec` objects that
        represents the type specification of elements of the iterator.

Raises:
        ValueError: If `dataset` is not provided and either `components` or
        `element_spec` is not provided. Or `dataset` is provided and either
        `components` and `element_spec` is provided.


Element-wise exponential.

Args:
        x: Tensor or variable.

Returns:
        A tensor.


def random_shuffle(value, seed=None, name=None):
    seed1, seed2 = random_seed.get_seed(seed)
    return gen_random_ops.random_shuffle(value, seed=seed1, seed2=seed2, name=name)

Randomly shuffles a tensor along its first dimension.

    The tensor is shuffled along dimension 0, such that each `value[j]` is mapped
    to one and only one `output[i]`. For example, a mapping that might occur for a
    3x2 tensor is:

    ```python
    [[1, 2],       [[5, 6],
    [3, 4],  ==>   [1, 2],
    [5, 6]]        [3, 4]]
    ```

Args:
        value: A Tensor to be shuffled.
        seed: A Python integer. Used to create a random seed for the distribution.
        See
        `tf.random.set_seed`
        for behavior.
        name: A name for the operation (optional).

Returns:
        A tensor of same shape and type as `value`, shuffled along its first
        dimension.


def _is_device_list_single_worker(devices):
    specs = []
    for d in devices:
        name = d.name if isinstance(d, context.LogicalDevice) else d
        specs.append(tf_device.DeviceSpec.from_string(name))
    num_workers = len({(d.job, d.task, d.replica) for d in specs})
    all_local = all((d.job in (None, 'localhost') for d in specs))
    any_local = any((d.job in (None, 'localhost') for d in specs))
    if any_local and (not all_local):
        raise ValueError("Local device should have only 'localhost' in the job field in device string. E.g. 'job:localhost' in /job:localhost/replica:0/task:0/device:CPU:0Devices cannot have mixed list of device strings containing both localhost and other job types such as worker, ps etc. ")
    if num_workers == 1 and (not all_local):
        if any((d.task is None for d in specs)):
            raise ValueError("Remote device string must have task specified.E.g. 'task:0' in /job:worker/replica:0/task:0/device:CPU:0")
    return num_workers == 1

Checks whether the devices list is for single or multi-worker.

Args:
        devices: a list of device strings or tf.config.LogicalDevice objects, for
        either local or for remote devices.

Returns:
        a boolean indicating whether these device strings are for local or for
        remote.

Raises:
        ValueError: if device strings are not consistent.


def on_test_batch_end(self, batch, logs=None):

Called at the end of a batch in `evaluate` methods.

    Also called at the end of a validation batch in the `fit`
    methods, if validation data is provided.

    Subclasses should override for any actions to run.

    Note that if the `steps_per_execution` argument to `compile` in
    `tf.keras.Model` is set to `N`, this method will only be called every `N`
    batches.

Args:
        batch: Integer, index of batch within the current epoch.
        logs: Dict. Aggregated metric results up until this batch.


Make a set of tests to do reduce operation.

Args:
        reduce_op: TensorFlow reduce operation to test, i.e. `tf.reduce_mean`.
        min_value: min value for created tensor data.
        max_value: max value for created tensor data.
        boolean_tensor_only: If true, will only generate tensor with boolean value.
        allow_fully_quantize: bool, whether fully_quantize is allowed.

Returns:
        a function representing the true generator with `reduce_op_in` curried.


def forward(self, input_ids: Optional[torch.LongTensor]=None, lang_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.FloatTensor]=None, token_type_ids: Optional[torch.LongTensor]=None, position_ids: Optional[torch.LongTensor]=None, head_mask: Optional[torch.FloatTensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, encoder_hidden_states: Optional[torch.FloatTensor]=None, encoder_attention_mask: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **kwargs) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    if labels is not None:
        use_cache = False
    outputs = self.roberta(input_ids, lang_ids=lang_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    sequence_output = outputs[0]
    prediction_scores = self.lm_head(sequence_output)
    lm_loss = None
    if labels is not None:
        lm_loss = self.loss_function(prediction_scores, labels, vocab_size=self.config.vocab_size, **kwargs)
    if not return_dict:
        output = (prediction_scores,) + outputs[2:]
        return (lm_loss,) + output if lm_loss is not None else output
    return CausalLMOutputWithCrossAttentions(loss=lm_loss, logits=prediction_scores, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, cross_attentions=outputs.cross_attentions)

lang_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
    Indices of the language adapters that should be activated for each sample, respectively. Default: the index
    that corresponds to `self.config.default_language`.
    labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
    Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
    `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
    ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`

Example:
    ```python
    >>> from transformers import AutoTokenizer, XmodForCausalLM, AutoConfig
    >>> import torch

    >>> tokenizer = AutoTokenizer.from_pretrained("FacebookAI/xlm-roberta-base")
    >>> config = AutoConfig.from_pretrained("facebook/xmod-base")
    >>> config.is_decoder = True
    >>> model = XmodForCausalLM.from_pretrained("facebook/xmod-base", config=config)
    >>> model.set_default_language("en_XX")

    >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
    >>> outputs = model(**inputs)

    >>> prediction_logits = outputs.logits
    ```


def forward(self, x):
    embeddings = self.embedding_convPxP(x).flatten(2)
    embeddings = nn.functional.pad(embeddings, (1, 0))
    embeddings = embeddings.permute(0, 2, 1)
    batch_size, sequence_length, embedding_dim = embeddings.shape
    embeddings = embeddings + self.positional_encoding_1d(batch_size, sequence_length, embedding_dim, device=embeddings.device, dtype=embeddings.dtype)
    for i in range(4):
        embeddings = self.transformer_encoder[i](embeddings)
    return embeddings

Forward pass

Args:
        x (torch.Tensor - NCHW): Input feature tensor

Returns:
        torch.Tensor - Transformer output embeddings of shape (batch_size, sequence_length, embedding_dim)


def __init__(self, config: XLMConfig):
    super().__init__()
    self.summary_type = getattr(config, 'summary_type', 'last')
    if self.summary_type == 'attn':
        raise NotImplementedError
    self.summary = nn.Identity()
    if hasattr(config, 'summary_use_proj') and config.summary_use_proj:
        if hasattr(config, 'summary_proj_to_labels') and config.summary_proj_to_labels and (config.num_labels > 0):
            num_classes = config.num_labels
        else:
            num_classes = config.hidden_size
        self.summary = nn.Linear(config.hidden_size, num_classes)
    activation_string = getattr(config, 'summary_activation', None)
    self.activation: Callable = get_activation(activation_string) if activation_string else nn.Identity()
    self.first_dropout = nn.Identity()
    if hasattr(config, 'summary_first_dropout') and config.summary_first_dropout > 0:
        self.first_dropout = nn.Dropout(config.summary_first_dropout)
    self.last_dropout = nn.Identity()
    if hasattr(config, 'summary_last_dropout') and config.summary_last_dropout > 0:
        self.last_dropout = nn.Dropout(config.summary_last_dropout)

Compute a single vector summary of a sequence hidden states.

Args:
        config ([`XLMConfig`]):
            The config used by the model. Relevant arguments in the config class of the model are (refer to the actual
            config class of your model for the default values it uses):

            - **summary_type** (`str`) -- The method to use to make this summary. Accepted values are:

                - `"last"` -- Take the last token hidden state (like XLNet)
                - `"first"` -- Take the first token hidden state (like Bert)
                - `"mean"` -- Take the mean of all tokens hidden states
                - `"cls_index"` -- Supply a Tensor of classification token position (GPT/GPT-2)
                - `"attn"` -- Not implemented now, use multi-head attention

            - **summary_use_proj** (`bool`) -- Add a projection after the vector extraction.
            - **summary_proj_to_labels** (`bool`) -- If `True`, the projection outputs to `config.num_labels` classes
            (otherwise to `config.hidden_size`).
            - **summary_activation** (`Optional[str]`) -- Set to `"tanh"` to add a tanh activation to the output,
            another string or `None` will add no activation.
            - **summary_first_dropout** (`float`) -- Optional dropout probability before the projection and activation.
            - **summary_last_dropout** (`float`)-- Optional dropout probability after the projection and activation.


def debug_string(self, indent: int=0) -> str:
    self_repr = f'{'| ' * indent}{self.__class__.__name__}<{repr(self)}>'
    if self._children:
        child_repr = '\n'.join((child.debug_string(indent=indent + 1) for child in self._children))
        self_repr = f'{self_repr}\n{child_repr}'
    return self_repr

Returns a debug string for the tree rooted at this node.

Args:
        indent: The level of indentation to begin printing the tree.

Returns:
        A string representing this node and its descendant nodes.


def hash_file(fpath, algorithm='sha256', chunk_size=65535):
    if isinstance(algorithm, str):
        hasher = resolve_hasher(algorithm)
    else:
        hasher = algorithm
    with open(fpath, 'rb') as fpath_file:
        for chunk in iter(lambda: fpath_file.read(chunk_size), b''):
            hasher.update(chunk)
    return hasher.hexdigest()

Calculates a file sha256 or md5 hash.

Example:
    >>> hash_file('/path/to/file.zip')
    'e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'

Args:
        fpath: Path to the file being validated.
        algorithm: Hash algorithm, one of `"auto"`, `"sha256"`, or `"md5"`.
            The default `"auto"` detects the hash algorithm in use.
        chunk_size: Bytes to read at a time, important for large files.

Returns:
        The file hash.


def _validate_or_infer_batch_size(self, batch_size, steps, x):
    if isinstance(x, (data_types.DatasetV1, data_types.DatasetV2, data_utils.Sequence)) or tf_inspect.isgenerator(x):
        if batch_size is not None:
            raise ValueError('The `batch_size` argument must not be specified for the given input type. Received input: {}, batch_size: {}'.format(x, batch_size))
        return
    layers = self._flatten_layers(include_self=False, recursive=False)
    first_layer = next(layers, None)
    if first_layer:
        static_batch_size = training_utils.get_static_batch_size(first_layer)
        if static_batch_size is not None:
            if self._distribution_strategy and distributed_training_utils.global_batch_size_supported(self._distribution_strategy):
                num_splits_for_ds = self._distribution_strategy.num_replicas_in_sync
            else:
                num_splits_for_ds = 1
            if batch_size is not None:
                if batch_size % num_splits_for_ds != 0:
                    raise ValueError('The `batch_size` argument ({}) must be divisible the by number of replicas ({})'.format(batch_size, num_splits_for_ds))
                per_replica_batch_size = batch_size // num_splits_for_ds
                if per_replica_batch_size != static_batch_size:
                    raise ValueError('The `batch_size` argument value {} is incompatible with the specified batch size of your Input Layer: {}'.format(per_replica_batch_size, static_batch_size))
            if isinstance(x, (data_types.DatasetV2, iterator_ops.Iterator, iterator_ops.IteratorBase)):
                ds_batch_size = tensor_shape.Dimension(nest.flatten(dataset_ops.get_legacy_output_shapes(x))[0][0]).value
                if ds_batch_size is not None:
                    if ds_batch_size % num_splits_for_ds != 0:
                        raise ValueError('The batch output shape of your `Dataset` {} cannot be divisible by number of replicas {}'.format(ds_batch_size, num_splits_for_ds))
                    ds_per_replica_batch_size = ds_batch_size // num_splits_for_ds
                    if ds_per_replica_batch_size != static_batch_size:
                        raise ValueError('The batch output shape of your `Dataset` is {}, which is incompatible with the specified batch size of your Input Layer: {}'.format(ds_per_replica_batch_size, static_batch_size))
            if steps is None:
                batch_size = static_batch_size * num_splits_for_ds
    if batch_size is None and steps is None:
        batch_size = 32
    return batch_size

Validates that the `batch_size` provided is consistent with InputLayer.

    It's possible that the user specified a static batch size in their
    InputLayer. If so, this method checks the provided `batch_size` and `x`
    arguments are consistent with this static batch size. Also, if
    `batch_size` is `None`, this method will attempt to infer the batch size
    from the static batch size of the InputLayer. Lastly, ValueError will be
    raised if `x` is a tf.data.Dataset and `batch_size` is specified as we
    expect users to provide batched datasets.

Args:
        batch_size: The batch_size provided as an argument to
        fit/evaluate/predict.
        steps: The steps provided as an argument to fit/evaluate/predict.
        x: The data passed as `x` to fit/evaluate/predict.

Returns:
        The validated batch_size, auto-inferred from the first layer if not
        provided.


def input_shape(self):
    return nest.map_structure(backend.int_shape, self.input)

Retrieves the input shape(s) of a layer.

    Only applicable if the layer has exactly one input,
    i.e. if it is connected to one incoming layer, or if all inputs
    have the same shape.

Returns:
        Input shape, as an integer shape tuple
        (or list of shape tuples, one tuple per input tensor).

Raises:
        AttributeError: if the layer has no defined input_shape.
        RuntimeError: if called in Eager mode.


def build_shuffle_all_reduce(input_tensors, gather_devices, red_op, un_op=None):
    input_tensors, shape = _flatten_tensors(input_tensors)
    dst_devices = [t.device for t in input_tensors]
    reduced_shards = _build_shuffle_gather(input_tensors, gather_devices, red_op, un_op)
    output_tensors = _build_shuffle_scatter(reduced_shards, dst_devices)
    if len(shape) != 1:
        output_tensors = _reshape_tensors(output_tensors, shape)
    return output_tensors

Construct a subgraph for shuffle all-reduce.

    Shuffle reduce is essentially the algorithm implemented when using
    parameter servers.  Suppose tensor length is n, there are d devices
    and g gather shards.  Each device sends a n/g length sub-tensor to
    each gather shard.  The gather shards perform a reduction across d
    fragments, then broadcast the result back to each device.  The
    devices then join the g fully reduced fragments they receive from
    the shards.  The gather shards could perform d-1 pairwise
    reductions, or one d-way reduction.  The first is better where
    reduction Op time is low compared to transmission time, the second
    better in the other case.

Args:
        input_tensors: list of `tf.Tensor` values to be reduced.
        gather_devices: list of names of devices on which reduction shards
        should be placed.
        red_op: an n-array elementwise reduction Op
        un_op: optional elementwise unary Op to be applied to fully-reduced values.

Returns:
        list of `tf.Tensor` which are the fully reduced tensors.


def _find_children_hints_in_while_loop(function_def, nodes_mapping):
    new_nodes = []
    for node in function_def.node_def:
        for i, _ in enumerate(node.input):
            if node.input[i] in nodes_mapping:
                node.input[i] = nodes_mapping[node.input[i]]
        new_nodes.append(_copy.deepcopy(node))
    name_to_seq_num = _extract_topology_sequence_mapping(function_def.node_def)
    children_hints = _find_all_hints_in_nodes(new_nodes)
    children_hints_q = []
    for hint in children_hints.values():
        _, output_names = hint.flattened_inputs_and_outputs()
        seq = name_to_seq_num[output_names[0]]
        for output_name in output_names:
            seq = min(seq, name_to_seq_num[output_name])
        children_hints_q.append((seq, hint))
    children_hints_q.sort(key=lambda tup: tup[0])
    ordered_children_hints = [x[1] for x in children_hints_q]
    return (ordered_children_hints, new_nodes)

Find children hints and all nodes inside the while loop.

Args:
        function_def: Function def of the while loop.
        nodes_mapping: While loop input_arg : real node name.

Returns:
        Ordered children hints and all re-mapped nodes inside the while loop.