def _psum(tensor, axis_name=None):
    if axis_name != _pmap_config.axis_name():
        raise ValueError('axis_name (%s) is not equal to that of the surrounding pmap (%s)' % (axis_name, _pmap_config.axis_name()))
    devices = _pmap_config.devices()
    if devices is None:
        raise ValueError("Can't retrieve the device list from the surrounding pmap")
    tensor = tf_np.asarray(tensor)
    if tpu_devices(devices):
        is_int64 = False
        is_float64 = False
        if tensor.dtype == np.int64:
            is_int64 = True
            tensor = tensor.astype(np.int32)
        elif tensor.dtype == np.float64:
            is_float64 = True
            tensor = tensor.astype(np.float32)
        tensor = tpu_ops.cross_replica_sum(tensor)
        if is_int64:
            tensor = math_ops.cast(tensor, dtypes.int64)
        elif is_float64:
            tensor = math_ops.cast(tensor, dtypes.float64)
    else:
        tensor = gen_collective_ops.collective_reduce(input=tensor, group_size=len(devices), group_key=_GROUP_KEY, instance_key=_get_instance_key(), merge_op='Add', final_op='Id', subdiv_offsets=(0,))
    return tf_np.asarray(tensor)

Sum all-reduction.

Args:
        tensor: A tensor.
        axis_name: The axis name to reduce. Must equal to that of the surrounding
        pmap.

Returns:
        The sum of the `tensor` replicas on each participating devices.


def testConcreteFunctionFlatSignatureError(self, conc_args=(), conc_kwargs=None, call_args=(), call_kwargs=None, error='.*', exception=TypeError):
    conc_args = conc_args() if callable(conc_args) else conc_args
    conc_kwargs = conc_kwargs() if callable(conc_kwargs) else conc_kwargs or {}
    call_args = call_args() if callable(call_args) else call_args
    call_kwargs = call_kwargs() if callable(call_kwargs) else call_kwargs or {}
    self.assertIsInstance(conc_args, tuple)
    self.assertIsInstance(call_args, tuple)
    self.assertIsInstance(conc_kwargs, dict)
    self.assertIsInstance(call_kwargs, dict)

    @polymorphic_function.function
    def func(x, y=5, *varargs, **kwargs):
        del y, varargs, kwargs
        return x
    conc = func.get_concrete_function(*conc_args, **conc_kwargs)
    with self.assertRaisesRegex(exception, error):
        self.evaluate(conc._call_with_flat_signature(call_args, call_kwargs))

Tests for errors in the flat signature.

Args:
        conc_args: Positional arguments used for get_concrete_function.
        conc_kwargs: Keyword arguments used for get_concrete_function.
        call_args: Positional arguments used to call the function.
        call_kwargs: Keyword arguments used to call the function.
        error: Expected exception message.
        exception: Expected exception type.


def _read_fhir_package_npm(npm_file: BinaryIO) -> Iterator[Tuple[str, str]]:
    with tarfile.open(fileobj=npm_file, mode='r:gz') as f:
        for member in f.getmembers():
            if member.name.endswith('.json'):
                content = f.extractfile(member)
                if content is not None:
                    yield (member.name, content.read().decode('utf-8-sig'))
            else:
                logging.info('Skipping  entry: %s.', member.name)

Yields the file entries for JSON resources in `npm_file` and their contents.

Args:
        npm_file: The `.tar.gz` or `.tgz` file following NPM conventions that should
        be parsed.

Yields:
        A tuple of filename and the raw JSON contained in that file.


def get_sine_pos_embed(pos_tensor: torch.Tensor, num_pos_feats: int=128, temperature: int=10000, exchange_xy: bool=True) -> Tensor:
    scale = 2 * math.pi
    dim_t = torch.arange(num_pos_feats, dtype=torch.float32, device=pos_tensor.device)
    dim_t = temperature ** (2 * torch.div(dim_t, 2, rounding_mode='floor') / num_pos_feats)

    def sine_func(x: torch.Tensor):
        sin_x = x * scale / dim_t
        sin_x = torch.stack((sin_x[..., 0::2].sin(), sin_x[..., 1::2].cos()), dim=3).flatten(2)
        return sin_x
    pos_tensor = pos_tensor.split([1] * pos_tensor.shape[-1], dim=-1)
    position_embeddings = [sine_func(x) for x in pos_tensor]
    if exchange_xy:
        position_embeddings[0], position_embeddings[1] = (position_embeddings[1], position_embeddings[0])
    position_embeddings = torch.cat(position_embeddings, dim=-1)
    return position_embeddings

Generate sine position embeddings from a position tensor.

Args:
        pos_tensor (torch.Tensor):
            Tensor containing positions. Shape: [..., n].
        num_pos_feats (`int`, *optional*, defaults to 128):
            Projected shape for each float in the tensor.
        temperature (`int`, *optional*, defaults to 10000):
            Temperature in the sine/cosine function.
        exchange_xy (`bool`, *optional*, defaults to `True`):
            Exchange pos x and pos y. For example, input tensor is [x,y], the results will be [pos(y), pos(x)].

Returns:
        position_embeddings (torch.Tensor): shape: [..., n * hidden_size].


def _testCompareToExplicitDerivative(self, dtype):
    delta = 0.001
    np_dtype = dtype.as_numpy_dtype
    try:
        from scipy import differentiate
        from scipy import special
        alpha_val = np.logspace(-2, 3, dtype=np_dtype)
        alpha = constant_op.constant(alpha_val)
        sample = random_ops.random_gamma([], alpha, np_dtype(1.0), dtype=dtype, seed=12345)
        actual = gradients_impl.gradients(sample, alpha)[0]
        sample_val, actual_val = self.evaluate((sample, actual))
        u = special.gammainc(alpha_val, sample_val)
        expected_val = differentiate.derivative(special.gammaincinv, alpha_val, args=(u,), initial_step=delta * alpha_val, order=2, preserve_shape=True).df
        self.assertAllClose(actual_val, expected_val, rtol=0.001, atol=0.001)
    except ImportError as e:
        tf_logging.warn('Cannot use special functions in a test: %s' % str(e))

Compare to the explicit reparameterization derivative.

    Verifies that the computed derivative satisfies
    dsample / dalpha = d igammainv(alpha, u) / dalpha,
    where u = igamma(alpha, sample).

Args:
        dtype: TensorFlow dtype to perform the computations in.


def _parse_test_option_args(self, argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('--test-pipeline-options', type=str, action='store', help='only run tests providing service options')
    parser.add_argument('--not-use-test-runner-api', action='store_true', default=False, help='whether not to use test-runner-api')
    known, unused_argv = parser.parse_known_args(argv)
    test_pipeline_options = known.test_pipeline_options or TestPipeline.pytest_test_pipeline_options
    if self.is_integration_test and (not test_pipeline_options):
        raise SkipTest('IT is skipped because --test-pipeline-options is not specified')
    self.not_use_test_runner_api = known.not_use_test_runner_api
    return shlex.split(test_pipeline_options) if test_pipeline_options else []

Parse value of command line argument: --test-pipeline-options to get
    pipeline options.

Args:
        argv: An iterable of command line arguments to be used. If not specified
        then sys.argv will be used as input for parsing arguments.

Returns:
        An argument list of options that can be parsed by argparser or directly
        build a pipeline option.


def replace_batch_norm(model):
    for name, module in model.named_children():
        if isinstance(module, nn.BatchNorm2d):
            new_module = DetaFrozenBatchNorm2d(module.num_features)
            if not module.weight.device == torch.device('meta'):
                new_module.weight.data.copy_(module.weight)
                new_module.bias.data.copy_(module.bias)
                new_module.running_mean.data.copy_(module.running_mean)
                new_module.running_var.data.copy_(module.running_var)
            model._modules[name] = new_module
        if len(list(module.children())) > 0:
            replace_batch_norm(module)

Recursively replace all `torch.nn.BatchNorm2d` with `DetaFrozenBatchNorm2d`.

Args:
        model (torch.nn.Module):
            input model


def variable_axis_size_partitioner(max_shard_bytes, axis=0, bytes_per_string_element=16, max_shards=None):
    if max_shard_bytes < 1 or bytes_per_string_element < 1:
        raise ValueError(f'Both max_shard_bytes and bytes_per_string_element must be positive. Currently, max_shard_bytes is {max_shard_bytes} andbytes_per_string_element is {bytes_per_string_element}')
    if max_shards and max_shards < 1:
        raise ValueError('max_shards must be positive.')

    def _partitioner(shape, dtype):
        """Partitioner that partitions shards to have max_shard_bytes total size.

    Args:
      shape: A `TensorShape`.
      dtype: A `DType`.

    Returns:
      A tuple representing how much to slice each axis in shape.

    Raises:
      ValueError: If shape is not a fully defined `TensorShape` or dtype is not
        a `DType`.
    """
        if not isinstance(shape, tensor_shape.TensorShape):
            raise ValueError(f'shape is not a TensorShape: {shape}')
        if not shape.is_fully_defined():
            raise ValueError(f'shape is not fully defined: {shape}')
        dtype = dtypes.as_dtype(dtype)
        if dtype.base_dtype == dtypes.string:
            element_size = bytes_per_string_element
        else:
            element_size = dtype.size
        partitions = [1] * shape.ndims
        bytes_per_slice = 1.0 * (shape.num_elements() / shape.dims[axis].value) * element_size
        slices_per_shard = max(1, math.floor(max_shard_bytes / bytes_per_slice))
        axis_shards = int(math.ceil(1.0 * shape.dims[axis].value / slices_per_shard))
        if max_shards:
            axis_shards = min(max_shards, axis_shards)
        partitions[axis] = axis_shards
        return partitions
    return _partitioner

Get a partitioner for VariableScope to keep shards below `max_shard_bytes`.

    This partitioner will shard a Variable along one axis, attempting to keep
    the maximum shard size below `max_shard_bytes`.  In practice, this is not
    always possible when sharding along only one axis.  When this happens,
    this axis is sharded as much as possible (i.e., every dimension becomes
    a separate shard).

    If the partitioner hits the `max_shards` limit, then each shard may end up
    larger than `max_shard_bytes`. By default `max_shards` equals `None` and no
    limit on the number of shards is enforced.

    One reasonable value for `max_shard_bytes` is `(64 << 20) - 1`, or almost
    `64MB`, to keep below the protobuf byte limit.

Args:
        max_shard_bytes: The maximum size any given shard is allowed to be.
        axis: The axis to partition along.  Default: outermost axis.
        bytes_per_string_element: If the `Variable` is of type string, this provides
        an estimate of how large each scalar in the `Variable` is.
        max_shards: The maximum number of shards in int created taking precedence
        over `max_shard_bytes`.

Returns:
        A partition function usable as the `partitioner` argument to
        `variable_scope` and `get_variable`.

Raises:
        ValueError: If any of the byte counts are non-positive.


def auto_docstring(obj=None, *, custom_intro=None, custom_args=None, checkpoint=None):

    def auto_docstring_decorator(obj):
        if len(obj.__qualname__.split('.')) > 1:
            return auto_method_docstring(obj, custom_args=custom_args, custom_intro=custom_intro, checkpoint=checkpoint)
        else:
            return auto_class_docstring(obj, custom_args=custom_args, custom_intro=custom_intro, checkpoint=checkpoint)
    if obj:
        return auto_docstring_decorator(obj)
    return auto_docstring_decorator

Automatically generates docstrings for classes and methods in the Transformers library.

    This decorator can be used in the following forms:
    @auto_docstring
    def my_function(...):
        ...
    or
    @auto_docstring()
    def my_function(...):
        ...
    or
    @auto_docstring(custom_intro="Custom intro", ...)
    def my_function(...):
        ...

Args:
        custom_intro (str, optional): Custom introduction text to add to the docstring. This will replace the default
            introduction text generated by the decorator before the Args section.
        checkpoint (str, optional): Checkpoint name to use in the docstring. This should be automatically inferred from the
            model configuration class, but can be overridden if needed.


def segment_max(data, segment_ids, num_segments=None, sorted=False):
    _segment_reduce_validation(data, segment_ids)
    if any_symbolic_tensors((data,)):
        return SegmentMax(num_segments, sorted).symbolic_call(data, segment_ids)
    return backend.math.segment_max(data, segment_ids, num_segments=num_segments, sorted=sorted)

Computes the max of segments in a tensor.

Args:
        data: Input tensor.
        segment_ids: A N-D tensor containing segment indices for each
            element in `data`. data.shape[:len(segment_ids.shape)] should match.
        num_segments: An integer representing the total number of
            segments. If not specified, it is inferred from the maximum
            value in `segment_ids`.
        sorted: A boolean indicating whether `segment_ids` is sorted.
            Defaults to `False`.

Returns:
        A tensor containing the max of segments, where each element
        represents the max of the corresponding segment in `data`.

Example:
    >>> data = keras.ops.convert_to_tensor([1, 2, 10, 20, 100, 200])
    >>> segment_ids = keras.ops.convert_to_tensor([0, 0, 1, 1, 2, 2])
    >>> num_segments = 3
    >>> keras.ops.segment_max(data, segment_ids, num_segments)
    array([2, 20, 200], dtype=int32)


def _set_state_variables(self, updates):
    if not self.built:
        raise RuntimeError('_set_state_variables() must be called after build().')
    with ops.init_scope():
        for var_name, value in updates.items():
            self.state_variables[var_name].assign(value)

Directly update the internal state of this Layer.

    This method expects a string-keyed dict of {state_variable_name: state}. The
    precise nature of the state, and the names associated, are describe by
    the subclasses of CombinerPreprocessingLayer.

Args:
        updates: A string keyed dict of weights to update.

Raises:
        RuntimeError: if 'build()' was not called before 'set_processing_state'.


def post_process_object_detection(self, outputs: 'Owlv2ObjectDetectionOutput', threshold: float=0.1, target_sizes: Optional[Union[TensorType, List[Tuple]]]=None):
    batch_logits, batch_boxes = (outputs.logits, outputs.pred_boxes)
    batch_size = len(batch_logits)
    if target_sizes is not None and len(target_sizes) != batch_size:
        raise ValueError('Make sure that you pass in as many target sizes as images')
    batch_class_logits = torch.max(batch_logits, dim=-1)
    batch_scores = torch.sigmoid(batch_class_logits.values)
    batch_labels = batch_class_logits.indices
    batch_boxes = center_to_corners_format(batch_boxes)
    if target_sizes is not None:
        batch_boxes = _scale_boxes(batch_boxes, target_sizes)
    results = []
    for scores, labels, boxes in zip(batch_scores, batch_labels, batch_boxes):
        keep = scores > threshold
        scores = scores[keep]
        labels = labels[keep]
        boxes = boxes[keep]
        results.append({'scores': scores, 'labels': labels, 'boxes': boxes})
    return results

Converts the raw output of [`Owlv2ForObjectDetection`] into final bounding boxes in (top_left_x, top_left_y,
    bottom_right_x, bottom_right_y) format.

Args:
        outputs ([`Owlv2ObjectDetectionOutput`]):
            Raw outputs of the model.
        threshold (`float`, *optional*, defaults to 0.1):
            Score threshold to keep object detection predictions.
        target_sizes (`torch.Tensor` or `List[Tuple[int, int]]`, *optional*):
            Tensor of shape `(batch_size, 2)` or list of tuples (`Tuple[int, int]`) containing the target size
            `(height, width)` of each image in the batch. If unset, predictions will not be resized.

Returns:
        `List[Dict]`: A list of dictionaries, each dictionary containing the following keys:
        - "scores": The confidence scores for each predicted box on the image.
        - "labels": Indexes of the classes predicted by the model on the image.
        - "boxes": Image bounding boxes in (top_left_x, top_left_y, bottom_right_x, bottom_right_y) format.


def _ragged_tensor_apply_loss(loss_fn, y_true, y_pred, y_pred_extra_dim=False):

    def rt_is_equiv_dense(rt):
        """Returns true if this RaggedTensor has the same row_lengths across

       all ragged dimensions and thus can be converted to a dense tensor
       without loss of information.

    Args:
      rt: RaggedTensor.
    """
        return math_ops.reduce_all([math_ops.equal(math_ops.reduce_variance(math_ops.cast(row_lens, backend.floatx())), constant_op.constant([0.0])) for row_lens in rt.nested_row_lengths()])

    def _convert_to_dense(inputs):
        return tuple((rt.to_tensor() if isinstance(rt, ragged_tensor.RaggedTensor) else rt for rt in inputs))

    def _call_loss(inputs, ragged_output):
        """ Adapt the result to ragged or dense tensor according to the expected

        output type. This is done so that all the return values of the map
        operation have the same type.
    """
        r = loss_fn(*inputs)
        if ragged_output and (not isinstance(r, ragged_tensor.RaggedTensor)):
            r = ragged_tensor.RaggedTensor.from_tensor(r)
        elif not ragged_output and isinstance(r, ragged_tensor.RaggedTensor):
            r = r.to_tensor()
        return r

    def _wrapper(inputs, ragged_output):
        _, y_pred = inputs
        if isinstance(y_pred, ragged_tensor.RaggedTensor):
            return cond.cond(rt_is_equiv_dense(y_pred), lambda: _call_loss(_convert_to_dense(inputs), ragged_output), lambda: _call_loss(inputs, ragged_output))
        return loss_fn(*inputs)
    if not isinstance(y_true, ragged_tensor.RaggedTensor):
        return loss_fn(y_true, y_pred.to_tensor())
    lshape = y_pred.shape.as_list()[1:-1]
    if len(lshape) > 0:
        spec = ragged_tensor.RaggedTensorSpec(shape=lshape, dtype=y_pred.dtype)
    else:
        spec = tensor_spec.TensorSpec(shape=[], dtype=y_pred.dtype)
    nested_splits_list = [rt.nested_row_splits for rt in (y_true, y_pred)]
    if y_pred_extra_dim:
        rdims = [len(slist) for slist in nested_splits_list]
        if rdims[0] == rdims[1] - 1:
            nested_splits_list[1] = nested_splits_list[1][:-1]
    map_fn = functools.partial(_wrapper, ragged_output=len(lshape) > 1)
    assertion_list = ragged_util.assert_splits_match(nested_splits_list)
    with ops.control_dependencies(assertion_list):
        return ragged_map_ops.map_fn(map_fn, elems=(y_true, y_pred), dtype=spec)

Apply a loss function on a per batch basis.

Args:
        loss_fn: The loss function
        y_true: truth values (RaggedTensor)
        y_pred: predicted values (RaggedTensor)
        y_pred_extra_dim: whether y_pred has an additional dimension compared to
        y_true

Returns:
        Loss-function result. A dense tensor if the output has a single dimension
        (per-batch loss value); a ragged tensor otherwise.


def __init__(self, min_length: int, eos_token_id: int):
    if not isinstance(min_length, int) or min_length < 0:
        raise ValueError(f'`min_length` has to be a positive integer, but is {min_length}')
    if not isinstance(eos_token_id, int) or eos_token_id < 0:
        raise ValueError(f'`eos_token_id` has to be a positive integer, but is {eos_token_id}')
    self.min_length = min_length
    self.eos_token_id = eos_token_id

[`FlaxLogitsProcessor`] enforcing a min-length by setting EOS probability to 0.

Args:
        min_length (`int`):
            The minimum length below which the score of `eos_token_id` is set to `-float("Inf")`.
        eos_token_id (`int`):
            The id of the *end-of-sequence* token.


Publicly accessible method for determining the current backend.

Returns:
        String, the name of the backend Keras is currently using. One of
            `"tensorflow"`, `"torch"`, or `"jax"`.

Example:
    >>> keras.config.backend()
    'tensorflow'


def get_replacement_transform_for_applied_ptransform(self, applied_ptransform):
    return self.get_replacement_transform(applied_ptransform.transform)

Provides a runner specific override for a given `AppliedPTransform`.

Args:
        applied_ptransform: `AppliedPTransform` containing the `PTransform` to be
        replaced.

Returns:
        A `PTransform` that will be the replacement for the `PTransform` inside
        the `AppliedPTransform` given as an argument.


def advise(self, options):
    advise_pb = tfprof_output_pb2.AdviceProto()
    opts = _build_advisor_options(options)
    advise_pb.ParseFromString(print_mdl.Profile('advise'.encode('utf-8'), opts.SerializeToString()))
    return advise_pb

Automatically detect problems and generate reports.

Args:
        options: A dict of options. See ALL_ADVICE example above.

Returns:
        An Advise proto that contains the reports from all checkers.


def get_device(ads, **kwargs):
    filtered = get_devices(ads, **kwargs)
    if len(filtered) == 1:
        return filtered[0]
    else:
        serials = [ad.serial for ad in filtered]
        raise Error('More than one device matched: %s' % serials)

Finds a unique AndroidDevice instance from a list that has specific
    attributes of certain values.

Example:
        get_device(android_devices, label='foo', phone_number='1234567890')
        get_device(android_devices, model='angler')

Args:
        ads: A list of AndroidDevice instances.
        kwargs: keyword arguments used to filter AndroidDevice instances.

Returns:
        The target AndroidDevice instance.

Raises:
        Error: None or more than one device is matched.


def forward(self, input_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, labels: Optional[torch.LongTensor]=None, next_sentence_label: Optional[torch.LongTensor]=None, output_attentions: Optional[torch.FloatTensor]=None, output_hidden_states: Optional[torch.FloatTensor]=None, return_dict: Optional[torch.FloatTensor]=None) -> Union[Tuple, MobileBertForPreTrainingOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    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)
    sequence_output, pooled_output = outputs[:2]
    prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)
    total_loss = None
    if labels is not None and next_sentence_label is not None:
        loss_fct = CrossEntropyLoss()
        masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
        next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
        total_loss = masked_lm_loss + next_sentence_loss
    if not return_dict:
        output = (prediction_scores, seq_relationship_score) + outputs[2:]
        return (total_loss,) + output if total_loss is not None else output
    return MobileBertForPreTrainingOutput(loss=total_loss, prediction_logits=prediction_scores, seq_relationship_logits=seq_relationship_score, hidden_states=outputs.hidden_states, attentions=outputs.attentions)

labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
    Labels for computing the masked language modeling loss. Indices should 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]`
    next_sentence_label (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
    Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair
    (see `input_ids` docstring) Indices should be in `[0, 1]`:

    - 0 indicates sequence B is a continuation of sequence A,
    - 1 indicates sequence B is a random sequence.

Example:
    ```python
    >>> from transformers import AutoTokenizer, MobileBertForPreTraining
    >>> import torch

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

    >>> input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)
    >>> # Batch size 1
    >>> outputs = model(input_ids)

    >>> prediction_logits = outputs.prediction_logits
    >>> seq_relationship_logits = outputs.seq_relationship_logits
    ```


def check_models_are_tested(module: types.ModuleType, test_file: str) -> List[str]:
    defined_models = get_models(module)
    tested_models = find_tested_models(test_file)
    if tested_models is None:
        if test_file.replace(os.path.sep, '/') in TEST_FILES_WITH_NO_COMMON_TESTS:
            return
        return [f'{test_file} should define `all_model_classes` to apply common tests to the models it tests. ' + 'If this intentional, add the test filename to `TEST_FILES_WITH_NO_COMMON_TESTS` in the file ' + '`utils/check_repo.py`.']
    failures = []
    for model_name, _ in defined_models:
        if model_name not in tested_models and should_be_tested(model_name):
            failures.append(f'{model_name} is defined in {module.__name__} but is not tested in ' + f'{os.path.join(PATH_TO_TESTS, test_file)}. Add it to the all_model_classes in that file.' + 'If common tests should not applied to that model, add its name to `IGNORE_NON_TESTED`' + 'in the file `utils/check_repo.py`.')
    return failures

Check models defined in a module are all tested in a given file.

Args:
        module (`types.ModuleType`): The module in which we get the models.
        test_file (`str`): The path to the file where the module is tested.

Returns:
        `List[str]`: The list of error messages corresponding to models not tested.


def meshgrid(*args, **kwargs):
    indexing = kwargs.pop('indexing', 'xy')
    name = kwargs.pop('name', 'meshgrid')
    if kwargs:
        key = list(kwargs.keys())[0]
        raise TypeError("'{}' is an invalid keyword argument for this function".format(key))
    if indexing not in ('xy', 'ij'):
        raise ValueError(f"Argument `indexing` parameter must be either 'xy' or 'ij', got '{indexing}'")
    with ops.name_scope(name, 'meshgrid', args) as name:
        ndim = len(args)
        s0 = (1,) * ndim
        if not ndim:
            return []
        output = []
        for i, x in enumerate(args):
            output.append(reshape(array_ops_stack.stack(x), s0[:i] + (-1,) + s0[i + 1:]))
        shapes = [size(x) for x in args]
        output_dtype = ops.convert_to_tensor(args[0]).dtype.base_dtype
        if indexing == 'xy' and ndim > 1:
            output[0] = reshape(output[0], (1, -1) + (1,) * (ndim - 2))
            output[1] = reshape(output[1], (-1, 1) + (1,) * (ndim - 2))
            shapes[0], shapes[1] = (shapes[1], shapes[0])
        mult_fact = ones(shapes, output_dtype)
        return [x * mult_fact for x in output]

Broadcasts parameters for evaluation on an N-D grid.

    Given N one-dimensional coordinate arrays `*args`, returns a list `outputs`
    of N-D coordinate arrays for evaluating expressions on an N-D grid.

Note:
    `meshgrid` supports cartesian ('xy') and matrix ('ij') indexing conventions.
    When the `indexing` argument is set to 'xy' (the default), the broadcasting
    instructions for the first two dimensions are swapped.

Example:
    Calling `X, Y = meshgrid(x, y)` with the tensors

    ```python
    x = [1, 2, 3]
    y = [4, 5, 6]
    X, Y = tf.meshgrid(x, y)
    # X = [[1, 2, 3],
    #      [1, 2, 3],
    #      [1, 2, 3]]
    # Y = [[4, 4, 4],
    #      [5, 5, 5],
    #      [6, 6, 6]]
    ```

Args:
        *args: `Tensor`s with rank 1.
        **kwargs:
        - indexing: Either 'xy' or 'ij' (optional, default: 'xy').
        - name: A name for the operation (optional).

Returns:
        outputs: A list of N `Tensor`s with rank N.

Raises:
        TypeError: When no keyword arguments (kwargs) are passed.
        ValueError: When indexing keyword argument is not one of `xy` or `ij`.


def _add_weight(self, name, initial_value, dtype=None):
    variable = variable_v1.VariableV1(initial_value=initial_value, name=name, dtype=dtype, trainable=False, use_resource=True, synchronization=variables.VariableSynchronization.AUTO, aggregation=variables.VariableAggregation.NONE)
    if context.executing_eagerly():
        graph_key = None
    else:
        graph = ops.get_default_graph()
        graph_key = graph._graph_key
    key = (name, graph_key)
    if self._weights.get(key, None) is not None:
        raise RuntimeError('Duplicate variables detected. {}'.format(key))
    self._weights[key] = variable
    self._handle_deferred_dependencies(name=name, trackable=variable)
    return variable

Adds a weight to this loss scale.

Args:
        name: Variable name.
        initial_value: The variable's initial value.
        dtype: The type of the variable.

Returns:
        A variable.

Raises:
        RuntimeError: If a weight with `name` has already been added.


def read_saved_model(saved_model_dir):
    path_to_pbtxt = os.path.join(compat.as_bytes(saved_model_dir), compat.as_bytes(constants.SAVED_MODEL_FILENAME_PBTXT))
    path_to_pb = os.path.join(compat.as_bytes(saved_model_dir), compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB))
    if not file_io.file_exists(path_to_pbtxt) and (not file_io.file_exists(path_to_pb)):
        raise IOError('SavedModel file does not exist at: %s' % saved_model_dir)
    saved_model = saved_model_pb2.SavedModel()
    if file_io.file_exists(path_to_pb):
        with file_io.FileIO(path_to_pb, 'rb') as f:
            file_content = f.read()
        try:
            saved_model.ParseFromString(file_content)
            return saved_model
        except message.DecodeError as e:
            raise IOError('Cannot parse proto file %s: %s.' % (path_to_pb, str(e)))
    elif file_io.file_exists(path_to_pbtxt):
        with file_io.FileIO(path_to_pbtxt, 'rb') as f:
            file_content = f.read()
        try:
            text_format.Merge(file_content.decode('utf-8'), saved_model)
            return saved_model
        except text_format.ParseError as e:
            raise IOError('Cannot parse pbtxt file %s: %s.' % (path_to_pbtxt, str(e)))
    else:
        raise IOError('SavedModel file does not exist at: %s/{%s|%s}' % (saved_model_dir, constants.SAVED_MODEL_FILENAME_PBTXT, constants.SAVED_MODEL_FILENAME_PB))

Reads the saved_model.pb or saved_model.pbtxt file containing `SavedModel`.

Args:
        saved_model_dir: Directory containing the SavedModel file.

Returns:
        A `SavedModel` protocol buffer.

Raises:
        IOError: If the file does not exist, or cannot be successfully parsed.


def actual_actual_isda(*, start_date, end_date, schedule_info=None, dtype=None, name=None):
    del schedule_info
    with tf.name_scope(name or 'actual_actual_isda'):
        end_date = dt.convert_to_date_tensor(end_date)
        start_date = dt.convert_to_date_tensor(start_date)
        dtype = dtype or tf.float32
        days_in_leap_years, days_in_nonleap_years = du.days_in_leap_and_nonleap_years_between(start_date, end_date)
        days_in_leap_years = tf.cast(days_in_leap_years, dtype=dtype)
        days_in_nonleap_years = tf.cast(days_in_nonleap_years, dtype=dtype)
        return days_in_leap_years / 366 + days_in_nonleap_years / 365

Computes the year fraction between the specified dates.

    Computes the year fraction between the dates by dividing the actual number of
    days in a leap year by 366 and the actual number of days in a standard year by
    365.

    When determining whether a leap day is contained in the date range,
    'start_date' is excluded and 'end_date' is included.

    Note that the schedule info is not needed for this convention and is ignored
    if supplied.

    https://en.wikipedia.org/wiki/Day_count_convention#Actual/Actual_ISDA

Args:
        start_date: A `DateTensor` object of any shape.
        end_date: A `DateTensor` object of compatible shape with `start_date`.
        schedule_info: The schedule info. Ignored for this convention.
        dtype: The dtype of the result. Either `tf.float32` or `tf.float64`. If not
        supplied, `tf.float32` is returned.
        name: Python `str` name prefixed to ops created by this function. If not
        supplied, `actual_actual_isda` is used.

Returns:
        A real `Tensor` of supplied `dtype` and shape of `start_date`. The year
        fraction between the start and end date as computed by Actual/Actual ISDA
        convention.


def __init__(self, cardinalities: List[int], embedding_dims: List[int]) -> None:
    super().__init__()
    self.num_features = len(cardinalities)
    self.embedders = nn.ModuleList([nn.Embedding(c, d) for c, d in zip(cardinalities, embedding_dims)])

Embed a sequence of categorical features.

Args:
        cardinalities (`list[int]`):
            List of cardinalities of the categorical features.
        embedding_dims (`list[int]`):
            List of embedding dimensions of the categorical features.


def __init__(self, key_dtype, value_dtype):
    self._key_dtype = dtypes.as_dtype(key_dtype)
    self._value_dtype = dtypes.as_dtype(value_dtype)

Construct a table initializer object.

Args:
        key_dtype: Type of the table keys.
        value_dtype: Type of the table values.


def GetShadowMap(self, since=None):
    return ShadowUpdateGetter(self.conf).GetUpdates(source=self, search_base=self.conf['base'], search_filter=self.conf['filter'], search_scope=self.conf['scope'], since=since)

Return the shadow map from this source.

Args:
        since: Get data only changed since this timestamp (inclusive) or None
        for all data.

Returns:
        instance of ShadowMap


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.FloatTensor]=None, head_mask: Optional[torch.FloatTensor]=None, past_key_values: Optional[Union[List[torch.FloatTensor], Cache]]=None, inputs_embeds: Optional[torch.FloatTensor]=None, start_positions: Optional[torch.LongTensor]=None, end_positions: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, position_ids: Optional[torch.LongTensor]=None) -> Union[Tuple, QuestionAnsweringModelOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    transformer_outputs = self.model(input_ids, past_key_values=past_key_values, attention_mask=attention_mask, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    hidden_states = transformer_outputs[0]
    logits = self.qa_outputs(hidden_states)
    start_logits, end_logits = logits.split(1, dim=-1)
    start_logits = start_logits.squeeze(-1).contiguous()
    end_logits = end_logits.squeeze(-1).contiguous()
    total_loss = None
    if start_positions is not None and end_positions is not None:
        if len(start_positions.size()) > 1:
            start_positions = start_positions.squeeze(-1)
        if len(end_positions.size()) > 1:
            end_positions = end_positions.squeeze(-1)
        ignored_index = start_logits.size(1)
        start_positions = start_positions.clamp(0, ignored_index).to(logits.device)
        end_positions = end_positions.clamp(0, ignored_index).to(logits.device)
        loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
        start_loss = loss_fct(start_logits, start_positions)
        end_loss = loss_fct(end_logits, end_positions)
        total_loss = (start_loss + end_loss) / 2
    if not return_dict:
        output = (start_logits, end_logits) + transformer_outputs[2:]
        return (total_loss,) + output if total_loss is not None else output
    return QuestionAnsweringModelOutput(loss=total_loss, start_logits=start_logits, end_logits=end_logits, hidden_states=transformer_outputs.hidden_states, attentions=transformer_outputs.attentions)

Example:

    ```python
    >>> from transformers import AutoTokenizer, OPTForQuestionAnswering
    >>> import torch

    >>> torch.manual_seed(4)  # doctest: +IGNORE_RESULT
    >>> tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")

    >>> # note: we are loading a OPTForQuestionAnswering from the hub here,
    >>> # so the head will be randomly initialized, hence the predictions will be random
    >>> model = OPTForQuestionAnswering.from_pretrained("facebook/opt-350m")

    >>> question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"

    >>> inputs = tokenizer(question, text, return_tensors="pt")
    >>> with torch.no_grad():
    ...     outputs = model(**inputs)

    >>> answer_start_index = outputs.start_logits.argmax()
    >>> answer_end_index = outputs.end_logits.argmax()

    >>> answer_offset = len(tokenizer(question)[0])

    >>> predict_answer_tokens = inputs.input_ids[
    ...     0, answer_offset + answer_start_index : answer_offset + answer_end_index + 1
    ... ]
    >>> predicted = tokenizer.decode(predict_answer_tokens)
    >>> predicted
    ' a nice puppet'
    ```


def test_batch_sample_paths_2d(self, batch_rank):
    dtype = tf.float64
    mu = np.array([0.2, 0.7])
    a = np.array([[0.4, 0.1], [0.3, 0.2]])
    b = np.array([[0.33, -0.03], [0.21, 0.5]])

    def drift_fn(t, x):
        return mu * tf.sqrt(t) * tf.ones_like(x, dtype=t.dtype)

    def vol_fn(t, x):
        return (a * t + b) * tf.ones(x.shape.as_list() + [2], dtype=t.dtype)
    process = tff.models.GenericItoProcess(dim=2, drift_fn=drift_fn, volatility_fn=vol_fn, dtype=dtype)
    times = np.array([0.1, 0.21, 0.32, 0.43, 0.55])
    x0 = np.array([0.1, -1.1]) * np.ones([2] * batch_rank + [1, 2])
    times_grid = None
    time_step = 0.01
    num_samples = 10000
    normal_draws = None
    paths = self.evaluate(process.sample_paths(times, num_samples=num_samples, initial_state=x0, time_step=time_step, times_grid=times_grid, normal_draws=normal_draws, seed=12134))
    num_samples = 10000
    self.assertAllClose(list(paths.shape), [2] * batch_rank + [num_samples, 5, 2], atol=0)
    means = np.mean(paths, axis=batch_rank)
    times = np.reshape(times, [1] * batch_rank + [-1, 1])
    expected_means = np.reshape(x0, [2] * batch_rank + [1, 2]) + 2.0 / 3.0 * mu * np.power(times, 1.5)
    self.assertAllClose(means, expected_means, rtol=0.01, atol=0.01)

Tests path properties for a batch of 2-dimentional Ito process.

    We construct the following Ito processes.

    dX_1 = mu_1 sqrt(t) dt + s11 dW_1 + s12 dW_2
    dX_2 = mu_2 sqrt(t) dt + s21 dW_1 + s22 dW_2

    mu_1, mu_2 are constants.
    s_ij = a_ij t + b_ij

    For this process expected value at time t is (x_0)_i + 2/3 * mu_i * t^1.5.

Args:
        batch_rank: The rank of the batch of processes being simulated.


def __init__(self, modules: Sequence[RelativePositionBiasBase]):
    super().__init__()
    self.biases = nn.ModuleList(modules)

Class which sums up various computed biases.

Args:
        modules (Sequence[RelativePositionBiasBase]):
            List of relative bias modules.


def patch_on_type(src: symbolic.Symbolic, value_type: Union[Type[Any], Tuple[Type[Any], ...]], value: Any=None, value_fn: Optional[Callable[[Any], Any]]=None, skip_notification: Optional[bool]=None) -> Any:
    return _conditional_patch(src, lambda k, v, p: isinstance(v, value_type), value, value_fn, skip_notification)

Recursively patch values on matched types.

    Example::

        d = pg.Dict(a={'x': 1}, b=2)
        print(pg.patching.patch_on_type(d, int, value_fn=lambda x: x * 2))
        # {a={x=2}, b=4}

Args:
        src: symbolic value to patch.
        value_type: Value type to match.
        value: New value for field that satisfy `condition`.
        value_fn: Callable object that produces new value based on old value.
        If not None, `value` must be None.
        skip_notification: If True, `on_change` event will not be triggered for this
        operation. If None, the behavior is decided by `pg.notify_on_rebind`.
        Please see `symbolic.Symbolic.rebind` for details.

Returns:
        `src` after being patched.


def __init__(self, alpha=1.0, **kwargs):
    super(ELU, self).__init__(**kwargs)
    if alpha is None:
        raise ValueError('Alpha of an ELU layer cannot be None, requires a float. Got %s' % alpha)
    self.supports_masking = True
    self.alpha = backend.cast_to_floatx(alpha)

Exponential Linear Unit.

    It follows:

    ```
        f(x) =  alpha * (exp(x) - 1.) for x < 0
        f(x) = x for x >= 0
    ```

    Input shape:
        Arbitrary. Use the keyword argument `input_shape`
        (tuple of integers, does not include the samples axis)
        when using this layer as the first layer in a model.

    Output shape:
        Same shape as the input.

Args:
        alpha: Scale for the negative factor.


def recipe_dcm_to_bigquery(config, auth_read, auth_write, account, report_id, report_name, dataset, table, is_incremental_load):
    dcm(config, {'auth': auth_read, 'report': {'account': account, 'report_id': report_id, 'name': report_name}, 'out': {'bigquery': {'auth': auth_write, 'dataset': dataset, 'table': table, 'header': True, 'is_incremental_load': is_incremental_load}}})

Move existing CM report into a BigQuery table.

Args:
        auth_read (authentication) - Credentials used for reading data.
        auth_write (authentication) - Credentials used for writing data.
        account (integer) - CM network id.
        report_id (integer) - CM report id, empty if using name .
        report_name (string) - CM report name, empty if using id instead.
        dataset (string) - Dataset to be written to in BigQuery.
        table (string) - Table to be written to in BigQuery.
        is_incremental_load (boolean) - Clear data in destination table during this report's time period, then append report data to existing table.


def filter_tensor(value: types.IntTensor, bit_mask: types.IntTensor, bit_index: types.IntTensor) -> types.IntTensor:
    is_bit_set = tf.equal(tf.cast(1, value.dtype), tf.bitwise.bitwise_and(tf.bitwise.right_shift(bit_mask, bit_index), tf.cast(1, value.dtype)))
    return tf.where(is_bit_set, value, 0)

Filters an input `Tensor` based on bits set in a mask `Tensor`.

    #### Examples

    ```python
    import tensorflow as tf
    import tf_quant_finance as tff

    # Example: Filtering a given vector based on a mask.

    tff.math.qmc.utils.filter_tensor(
        tf.constant([5, 6, 7, 8]),
        tf.constant([1, 2, 3, 4]),
        tf.constant([1, 1, 2, 2]))
    # ==> tf.Tensor([0, 6, 0, 8], shape=(4,), dtype=int32)
    ```

Args:
        value: Scalar `Tensor` of integers.
        bit_mask: Positive scalar `Tensor` of integers with the same `shape` and
        `dtype` as `value`.
        bit_index: Positive scalar `Tensor` of integers with the same `shape` and
        `dtype` as `value`.

Returns:
        `Tensor` with the same `shape` as `value` equal to `value` if the
        `bit_index`-th bit in `LSB 0` order is set in `bit_mask`, or `zero`
        otherwise i.e.: `value * (1 & (bit_mask >>> bit_index))`


def __init__(self, primals, tangents):
    self._accumulator = pywrap_tfe.TFE_Py_ForwardAccumulatorNew(False)
    self._recording = False
    primal_ids = set()
    for primal in nest.flatten(primals):
        if id(primal) in primal_ids:
            raise ValueError('Tensor {} was specified as a primal multiple times. This may indicate an error. If it was intended, please sum the corresponding tangents.')
        primal_ids.add(id(primal))
    self._watch(primals, tangents)

Specify tensors to watch and their Jacobian-vector products.

    Mathematically, `tangents` is a vector right-multiplying the Jacobian matrix
    (a Jacobian-vector product) for the function computed while this accumulator
    is active. Since JVPs are computed in forward mode as the computation
    happens, this vector must be supplied in advance.

    Listing a single tensor multiple times in `primals` raises an
    exception. Excluding a tensor from `primals` is equivalent to watching it
    with a tangent tensor of zeros.

Args:
        primals: A tensor or nested structure of tensors to watch.
        tangents: A tensor or nested structure of tensors, with the same nesting
        structure as `primals`, with each element being a vector with the same
        size as the corresponding primal element.

Raises:
        ValueError: If the same tensor or variable is specified multiple times in
        `primals`.


def check(self, src, expected=None, prologue=None, name=None, version=None, platform='linux'):
    ast = self.parse(src, name, version, platform)
    actual = pytd_utils.Print(ast)
    if expected != IGNORE:
        if expected is None:
            expected = src
        expected = textwrap.dedent(expected).lstrip()
        if prologue:
            expected = f'{textwrap.dedent(prologue)}\n\n{expected}'
        self.assertMultiLineEqual(expected.rstrip(), actual)
    return ast

Check the parsing of src.

    This checks that parsing the source and then printing the resulting
    AST results in the expected text.

Args:
        src: A source string.
        expected: Optional expected result string.  If not provided, src is used
        instead.  The special value IGNORE can be used to skip checking the
        parsed results against expected text.
        prologue: An optional prologue to be prepended to the expected text before
        comparison.  Useful for imports that are introduced during printing the
        AST.
        name: The name of the module.
        version: A python version tuple (None for default value).
        platform: A platform string (defaults to "linux").

Returns:
        The parsed pytd.TypeDeclUnit.


def get_image(width, height, want_grayscale, filepath):
    with ops.Graph().as_default():
        with session.Session():
            file_data = io_ops.read_file(filepath)
            channels = 1 if want_grayscale else 3
            image_tensor = image_ops.decode_image(file_data, channels=channels).eval()
            resized_tensor = image_ops.resize_images_v2(image_tensor, (height, width)).eval()
    return resized_tensor

Returns an image loaded into an np.ndarray with dims [height, width, (3 or 1)].

Args:
        width: Width to rescale the image to.
        height: Height to rescale the image to.
        want_grayscale: Whether the result should be converted to grayscale.
        filepath: Path of the image file..

Returns:
        np.ndarray of shape (height, width, channels) where channels is 1 if
        want_grayscale is true, otherwise 3.


Checks whether the server is still running.

    If the server is not running, it throws an error. As this function is called
    each time the client tries to send an RPC, this should be a quick check
    without affecting performance. Otherwise it is fine to not check anything.

Raises:
        errors.ServerDiedError: if the server died.


def extensions(self, prefix: str):
    prefix_node = self._get_node(prefix)
    ret = self._collect_tokens(prefix_node)
    return [prefix + token for token in ret]

Generates all extensions of a given prefix token in the Trie.

Example:
    ```python
    >>> trie = Trie()
    >>> trie.add("apple")
    >>> trie.add("app")
    >>> trie.add("application")
    >>> trie.extensions("app")
    ['app', 'apple', 'application']
    ```


def set_number_of_shards(self, number_of_shards):
    for policy in self._sharding_policies:
        policy.set_number_of_shards(number_of_shards)
        policy.set_number_of_partitions(self._number_of_partitions)
    self._validate()

Sets the number of shards to use for the InfeedQueue.

Args:
        number_of_shards: number of ways to shard the InfeedQueue.

Raises:
        ValueError: if number_of_shards is not > 0; or the policies have
        been frozen and number_of_shards was already set to something
        else.


def show_backref(target, max_depth=3):
    if objgraph is None:
        raise NotImplementedError('objgraph is not installed.')
    string_io = io.StringIO()
    objgraph.show_backrefs(target, max_depth=max_depth, output=string_io)
    graph = string_io.getvalue()
    string_io.close()
    return graph

Returns a dot graph of all the objects that are referencing the target.

    A object referencing graph is useful to debug memory leak like circular
    reference. objgraph provides a good visualization of the memory graph than
    most python built-in utilities like gc.get_referrers(), which are not
    human-readable sometimes.

    The dot graph will be written to a string IO object, and can be rendered with
    graphviz in operating system.
    E.g. dot -Tpng {$dot_graph} -o output.png

Args:
        target: The target object for the memory graph.
        max_depth: The maximum depth of the graph. By default 3 layers of references
        are used. Increases this a lot may result in the graph growing too big.

Returns:
        A string that contains the object reference graph.

Raises:
        NotImplementedError: if objgraph is not installed.


def get_auto_remove_tmp_dir(self, tmp_dir=None, before=None, after=None):
    if tmp_dir is not None:
        if before is None:
            before = True
        if after is None:
            after = False
        path = Path(tmp_dir).resolve()
        if not tmp_dir.startswith('./'):
            raise ValueError(f'`tmp_dir` can only be a relative path, i.e. `./some/path`, but received `{tmp_dir}`')
        if before is True and path.exists():
            shutil.rmtree(tmp_dir, ignore_errors=True)
        path.mkdir(parents=True, exist_ok=True)
    else:
        if before is None:
            before = True
        if after is None:
            after = True
        tmp_dir = tempfile.mkdtemp()
    if after is True:
        self.teardown_tmp_dirs.append(tmp_dir)
    return tmp_dir

Args:
    tmp_dir (`string`, *optional*):
        if `None`:

        - a unique temporary path will be created
        - sets `before=True` if `before` is `None`
        - sets `after=True` if `after` is `None`
        else:

        - `tmp_dir` will be created
        - sets `before=True` if `before` is `None`
        - sets `after=False` if `after` is `None`
    before (`bool`, *optional*):
        If `True` and the `tmp_dir` already exists, make sure to empty it right away if `False` and the
        `tmp_dir` already exists, any existing files will remain there.
    after (`bool`, *optional*):
        If `True`, delete the `tmp_dir` at the end of the test if `False`, leave the `tmp_dir` and its contents
        intact at the end of the test.

Returns:
    tmp_dir(`string`): either the same value as passed via *tmp_dir* or the path to the auto-selected tmp dir


def generate_typegraph(program: cfg.Program, var_table: dict[int, str], loader: jinja2.BaseLoader) -> str:
    encoder = typegraph_serializer.TypegraphEncoder()
    enc_prog = encoder.default(program)
    return _generate_visualization(template_file=_TYPEGRAPH_TEMPLATE_NAME, loader=loader, program=json.dumps(enc_prog), query_table=enc_prog['queries'], var_table=var_table)

Generate the visualization webpage.

Args:
        program: cfg.Program. The instance of the program to visualize.
        var_table: dict[int, str]. A mapping of cfg.Variable IDs to names.
        loader: A jinja2 loader

Returns:
        str. The rendered visualization page.


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

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

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

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

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


def _process_new(self, feed_item):
    lp = self.landing_page_dao.get(feed_item, required=True)
    feed_item[FieldMap.CAMPAIGN_LANDING_PAGE_ID] = lp['id']
    feed_item[FieldMap.CAMPAIGN_LANDING_PAGE_NAME] = lp['name']
    return {'advertiserId': feed_item.get(FieldMap.ADVERTISER_ID, None), 'name': feed_item.get(FieldMap.CAMPAIGN_NAME, None), 'startDate': StringExtensions.convertDateTimeStrToDateStr(feed_item.get(FieldMap.CAMPAIGN_START_DATE, None)), 'endDate': StringExtensions.convertDateTimeStrToDateStr(feed_item.get(FieldMap.CAMPAIGN_END_DATE, None)), 'defaultLandingPageId': lp['id']}

Creates a new campaign DCM object from a feed item representing a campaign from the Bulkdozer feed.

    This function simply creates the object to be inserted later by the BaseDAO
    object.

Args:
        feed_item: Feed item representing the campaign from the Bulkdozer feed.

Returns:
        A campaign object ready to be inserted in DCM through the API.


def __init__(self, name):
    stl.base.ParameterizedObject.__init__(self, name)
    self.expand = None

User defined events.

    A user defined event must be an expression of a external event with arguments.
    For example,

    event eUserDefinedEvent = eBuiltInEvent(arg1, arg2);

    It provides all information to turn event.EventInTransition into
    stl.base.Func.

Attributes:
        expand: Expression to expand to a resolved event, stl.base.Func.


def expand(self, pcoll: beam.PCollection[Union[beam.Row, NamedTuple]]) -> beam.PCollection[common_types.InstanceDictType]:
    return pcoll | beam.Map(lambda x: x._asdict())

Args:
    pcoll: A PCollection of NamedTuples or Rows.

Returns:
    A PCollection of dictionaries.


def pad_image(self, image: np.ndarray, size: Dict[str, int], random_padding: bool=False, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> np.ndarray:
    output_height, output_width = (size['height'], size['width'])
    input_height, input_width = get_image_size(image, channel_dim=input_data_format)
    delta_width = output_width - input_width
    delta_height = output_height - input_height
    if random_padding:
        pad_top = np.random.randint(low=0, high=delta_height + 1)
        pad_left = np.random.randint(low=0, high=delta_width + 1)
    else:
        pad_top = delta_height // 2
        pad_left = delta_width // 2
    pad_bottom = delta_height - pad_top
    pad_right = delta_width - pad_left
    padding = ((pad_top, pad_bottom), (pad_left, pad_right))
    return pad(image, padding, data_format=data_format, input_data_format=input_data_format)

Pad the image to the specified size.

Args:
        image (`np.ndarray`):
            The image to be padded.
        size (`Dict[str, int]`):
            The size `{"height": h, "width": w}` to pad the image to.
        random_padding (`bool`, *optional*, defaults to `False`):
            Whether to use random padding or not.
        data_format (`str` or `ChannelDimension`, *optional*):
            The data format of the output image. If unset, the same format as the input image is used.
        input_data_format (`ChannelDimension` or `str`, *optional*):
            The channel dimension format of the input image. If not provided, it will be inferred.


def __init__(self, bytes_per_pack=0, timeout_seconds=None):
    """Creates a CollectiveHints.

Args:
  bytes_per_pack: a non-negative integer. Breaks collective operations into
    packs of certain size. If it's zero, the value is determined
    automatically. This only applies to all-reduce with
    `MultiWorkerMirroredStrategy` currently.
  timeout_seconds: a float or None, timeout in seconds. If not None, the
    collective raises `tf.errors.DeadlineExceededError` if it takes longer
    than this timeout. This can be useful when debugging hanging issues.
    This should only be used for debugging since it creates a new thread for
    each collective, i.e. an overhead of `timeout_seconds *
    num_collectives_per_second` more threads.  This only works for
    `tf.distribute.experimental.MultiWorkerMirroredStrategy`.

Raises:
  ValueError: When arguments have invalid value.
"""
    pass

Hints for collective operations like AllReduce.

    This can be passed to methods like
    `tf.distribute.get_replica_context().all_reduce()` to optimize collective
    operation performance. Note that these are only hints, which may or may not
    change the actual behavior. Some options only apply to certain strategy and
    are ignored by others.

    One common optimization is to break gradients all-reduce into multiple packs
    so that weight updates can overlap with gradient all-reduce.

Example:
    - bytes_per_pack

    ```python
    hints = tf.distribute.experimental.CollectiveHints(
        bytes_per_pack=50 * 1024 * 1024)
    grads = tf.distribute.get_replica_context().all_reduce(
        'sum', grads, experimental_hints=hints)
    optimizer.apply_gradients(zip(grads, vars),
        experimental_aggregate_gradients=False)
    ```

    - timeout_seconds

    ```python
    strategy = tf.distribute.MirroredStrategy()
    hints = tf.distribute.experimental.CollectiveHints(
        timeout_seconds=120.0)
    try:
        strategy.reduce("sum", v, axis=None, experimental_hints=hints)
    except tf.errors.DeadlineExceededError:
        do_something()
    ```


def __init__(self, count_mode='samples', stateful_metrics=None):
    super(ProgbarLogger, self).__init__()
    self._supports_tf_logs = True
    if count_mode == 'samples':
        self.use_steps = False
    elif count_mode == 'steps':
        self.use_steps = True
    else:
        raise ValueError('Unknown `count_mode`: ' + str(count_mode))
    self.stateful_metrics = set(stateful_metrics) if stateful_metrics else set()
    self.seen = 0
    self.progbar = None
    self.target = None
    self.verbose = 1
    self.epochs = 1
    self._train_step, self._test_step, self._predict_step = (None, None, None)
    self._call_batch_hooks = True
    self._called_in_fit = False

Callback that prints metrics to stdout.

Args:
        count_mode: One of `"steps"` or `"samples"`.
            Whether the progress bar should
            count samples seen or steps (batches) seen.
        stateful_metrics: Iterable of string names of metrics that
            should *not* be averaged over an epoch.
            Metrics in this list will be logged as-is.
            All others will be averaged over time (e.g. loss, etc).
            If not provided, defaults to the `Model`'s metrics.

Raises:
        ValueError: In case of invalid `count_mode`.


def binomial(self, shape, counts, probs, dtype=dtypes.int32, name=None):
    dtype = dtypes.as_dtype(dtype)
    with ops.name_scope(name, 'binomial', [shape, counts, probs]) as name:
        counts = ops.convert_to_tensor(counts, name='counts')
        probs = ops.convert_to_tensor(probs, name='probs')
        shape_tensor = _shape_tensor(shape)
        return gen_stateful_random_ops.stateful_random_binomial(self.state.handle, self.algorithm, shape=shape_tensor, counts=counts, probs=probs, dtype=dtype, name=name)

Outputs random values from a binomial distribution.

    The generated values follow a binomial distribution with specified count and
    probability of success parameters.

Example:
    ```python
    counts = [10., 20.]
    # Probability of success.
    probs = [0.8]

    rng = tf.random.Generator.from_seed(seed=234)
    binomial_samples = rng.binomial(shape=[2], counts=counts, probs=probs)

    counts = ... # Shape [3, 1, 2]
    probs = ...  # Shape [1, 4, 2]
    shape = [3, 4, 3, 4, 2]
    rng = tf.random.Generator.from_seed(seed=1717)
    # Sample shape will be [3, 4, 3, 4, 2]
    binomial_samples = rng.binomial(shape=shape, counts=counts, probs=probs)
    ```

Args:
        shape: A 1-D integer Tensor or Python array. The shape of the output
        tensor.
        counts: Tensor. The counts of the binomial distribution. Must be
        broadcastable with `probs`, and broadcastable with the rightmost
        dimensions of `shape`.
        probs: Tensor. The probability of success for the
        binomial distribution. Must be broadcastable with `counts` and
        broadcastable with the rightmost dimensions of `shape`.
        dtype: The type of the output. Default: tf.int32
        name: A name for the operation (optional).

Returns:
        samples: A Tensor of the specified shape filled with random binomial
        values.  For each i, each samples[i, ...] is an independent draw from
        the binomial distribution on counts[i] trials with probability of
        success probs[i].


def __init__(self, input_resolution: Tuple[int], dim: int, norm_layer: nn.Module=nn.LayerNorm) -> None:
    super().__init__()
    self.input_resolution = input_resolution
    self.dim = dim
    self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
    self.norm = norm_layer(4 * dim)

Patch Merging Layer.

Args:
        input_resolution (`Tuple[int]`):
            Resolution of input feature.
        dim (`int`):
            Number of input channels.
        norm_layer (`nn.Module`, *optional*, defaults to `nn.LayerNorm`):
            Normalization layer class.


def make(self, path, metadata=None):
    self.current_shard_filenames = []
    if self.h5_file is not None:
        self.current_shard_filenames.append(pathlib.Path(self.h5_file.filename).name)
    return super().make(path, metadata)

Make a new H5 entry group.

    This method is only available in write mode. It defers the creation of
    the H5 entry group until `__setitem__` is called, preventing the
    creation of empty groups.

    The information about the current shard is reset.

Args:
        path: `str`. The variable path.
        metadata: Optional `dict`. The metadata to save with the H5 entry
            group. Defaults to `None`.


def check_termination(self) -> None:
    if self._is_thread_joined:
        if self.is_alive():
            raise RuntimeError('Thread was not joined with main thread, and is still running when the test finished.')
    else:
        self._testcase.fail('A checked thread was not joined.')

Returns whether the checked thread was properly used and did terminate.

    Every checked thread should be "join"ed after starting, and before the
    test tears down. If it is not joined, it is possible the thread will hang
    and cause flaky failures in tests.

Raises:
        self._testcase.failureException: If check_termination was called before
        thread was joined.

        RuntimeError: If the thread is not terminated. This means thread was not
        joined with the main thread.


def __init__(self, trainer_id):
    if not trainer_id:
        raise ValueError('tf.data service cross-trainer cache requires a non-empty trainer ID.')
    self.trainer_id = trainer_id

Constructs a CrossTrainerCache.

Args:
        trainer_id: Each training job has a unique ID. Once a job has consumed
        data, the data remains in the cache and is re-used by jobs with different
        `trainer_id`s. Requests with the same `trainer_id` do not re-use data.

Raises:
        ValueError if `trainer_id` is empty.


def npmfile_containing(file_contents: Sequence[Tuple[str, str]]):
    with tempfile.NamedTemporaryFile(suffix='.tgz') as temp_file:
        with tarfile.open(fileobj=temp_file, mode='w:gz') as tar_file:
            for file_name, contents in file_contents:
                info = tarfile.TarInfo(name=f'package/{file_name}')
                info.size = len(contents)
                tar_file.addfile(tarinfo=info, fileobj=io.BytesIO(contents.encode('utf-8')))
        temp_file.flush()
        yield temp_file

Builds a temp file containing a NPM .tar.gz file with the given contents.

Args:
        file_contents: Sequence of (file_name, file_contents) tuples to be written
        to the NPM file.

Yields:
        A tempfile.NamedTemporaryFile for the written zip file.


def _get_ops_from_nodedefs(node_defs):
    ops = set()
    for node_def in node_defs:
        op_and_kernel = get_ops_from_nodedef(node_def)
        if op_and_kernel:
            ops.add(op_and_kernel)
    return ops

Gets the ops and kernels needed from the list of NodeDef.

    If a NodeDef's op is not in the allowlist of ops without kernel and there is
    no kernel found for this NodeDef, then skip that NodeDef and proceed to the
    next one.

Args:
        node_defs: list of NodeDef's to get op/kernel information.

Returns:
        A set of (op_name, kernel_name) tuples.


def _all_reduce(self, reduce_op, value, replica_id, options):
    raise NotImplementedError('_all_reduce must be implemented in descendants.')

All-reduce the `value` across all replicas so that all get the result.

    `value` can be a nested structure of tensors or `IndexedSlices`. The
    implementation should generally batch the all-reduces when possible.
    `options` can be set to hint the batching behavior.

    This API must be called in a replica context.

Args:
        reduce_op: A `tf.distribute.ReduceOp` value specifying how values should
        be combined.
        value: Value to be reduced. A tensor or a nested structure of tensors or
        `IndexedSlices`.
        replica_id: An integer indicating the id of the replica where this
        all_reduce is called under. This is the local replica id that ranges
        from 0 to len(local_devices) - 1.
        options: A `tf.distribute.experimental.CommunicationOptions`.

Returns:
        A tensor/IndexedSlices or a nested structure of tensors/IndexedSlices with
        the reduced values. The structure is the same as `value`.


def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, layer_head_mask: torch.Tensor, output_attentions: bool=False, output_router_logits: bool=False) -> torch.Tensor:
    residual = hidden_states
    hidden_states = self.self_attn_layer_norm(hidden_states)
    hidden_states, attn_weights, _ = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask, layer_head_mask=layer_head_mask, output_attentions=output_attentions)
    hidden_states = self.attn_dropout(hidden_states)
    hidden_states = residual + hidden_states
    residual = hidden_states
    hidden_states = self.ff_layer_norm(hidden_states)
    if self.is_sparse:
        hidden_states, router_states = self.ffn(hidden_states, attention_mask)
    else:
        hidden_states, router_states = (self.ffn(hidden_states), None)
    hidden_states = self.ff_dropout(hidden_states)
    hidden_states = residual + hidden_states
    if hidden_states.dtype == torch.float16 and (torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()):
        clamp_value = torch.finfo(hidden_states.dtype).max - 1000
        hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
    outputs = (hidden_states,)
    if output_attentions:
        outputs += (attn_weights,)
    if output_router_logits:
        outputs += (router_states,)
    return outputs

Args:
    hidden_states (`torch.FloatTensor`):
        input to the layer of shape `(batch, seq_len, embed_dim)`
    attention_mask (`torch.FloatTensor`):
        attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very
        large negative values.
    layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size
        `(encoder_attention_heads,)`.
    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.


def get_fhir_extensions(msg: message.Message) -> List[message.Message]:
    extension_field = msg.DESCRIPTOR.fields_by_name.get('extension')
    if extension_field is None:
        raise ValueError(f'Message of type {msg.DESCRIPTOR.name} does not have an extension attribute.')
    return proto_utils.get_value_at_field(msg, extension_field)

Returns a list of FHIR extensions from the provided FHIR primitive.

Args:
        msg: The Message to extract FHIR extensions from.

Returns:
        A list of Extension values found on the instance.

Raises:
        ValueError: If the provided FHIR primitive doesn't have an extension
        attribute.


def _get_scope(node_name):
    if not node_name:
        raise ValueError(f'Node name cannot be empty or None. Received: {node_name}.')
    if node_name.startswith('^'):
        node_name = node_name[1:]
    if '/' in node_name:
        scope, _ = node_name.rsplit('/', 1)
        return scope
    return ''

Extract the scope name from a node name.

    The scope name is everything before the final slash,
    not including any ^ prefix denoting a control dependency.

Args:
        node_name: the full name of an Op or a Tensor in the graph.

Returns:
        The deepest named scope containing the node.

Raises:
        ValueError: if tensor_name is None or empty


def split(self, text: str) -> list[str]:
    states = OrderedDict()
    offsets = [0]
    skip = 0
    for current, current_char in enumerate(text):
        if skip and current < skip:
            continue
        to_remove = set()
        reset = False
        for start, trie_pointer in states.items():
            if '' in trie_pointer:
                for lookstart, looktrie_pointer in states.items():
                    if lookstart > start:
                        break
                    elif lookstart < start:
                        lookahead_index = current + 1
                        end = current + 1
                    else:
                        lookahead_index = current
                        end = current
                    next_char = text[lookahead_index] if lookahead_index < len(text) else None
                    if '' in looktrie_pointer:
                        start = lookstart
                        end = lookahead_index
                        skip = lookahead_index
                    while next_char in looktrie_pointer:
                        looktrie_pointer = looktrie_pointer[next_char]
                        lookahead_index += 1
                        if '' in looktrie_pointer:
                            start = lookstart
                            end = lookahead_index
                            skip = lookahead_index
                        if lookahead_index == len(text):
                            break
                        next_char = text[lookahead_index]
                offsets.append(start)
                offsets.append(end)
                reset = True
                break
            elif current_char in trie_pointer:
                trie_pointer = trie_pointer[current_char]
                states[start] = trie_pointer
            else:
                to_remove.add(start)
        if reset:
            states = {}
        else:
            for start in to_remove:
                del states[start]
        if current >= skip and current_char in self.data:
            states[current] = self.data[current_char]
    for start, trie_pointer in states.items():
        if '' in trie_pointer:
            end = len(text)
            offsets.append(start)
            offsets.append(end)
            break
    return self.cut_text(text, offsets)

Will look for the words added to the trie within `text`. Output is the original string splitted along the
    boundaries of the words found.

    This trie will match the longest possible word first !

Example:
    ```python
    >>> trie = Trie()
    >>> trie.split("[CLS] This is a extra_id_100")
    ["[CLS] This is a extra_id_100"]

    >>> trie.add("[CLS]")
    >>> trie.add("extra_id_1")
    >>> trie.add("extra_id_100")
    >>> trie.split("[CLS] This is a extra_id_100")
    ["[CLS]", " This is a ", "extra_id_100"]
    ```


def weighting_function(max_num_bins: int, up: torch.Tensor, reg_scale: int) -> torch.Tensor:
    upper_bound1 = abs(up[0]) * abs(reg_scale)
    upper_bound2 = abs(up[0]) * abs(reg_scale) * 2
    step = (upper_bound1 + 1) ** (2 / (max_num_bins - 2))
    left_values = [-step ** i + 1 for i in range(max_num_bins // 2 - 1, 0, -1)]
    right_values = [step ** i - 1 for i in range(1, max_num_bins // 2)]
    values = [-upper_bound2] + left_values + [torch.zeros_like(up[0][None])] + right_values + [upper_bound2]
    values = torch.cat(values, 0)
    return values

Generates the non-uniform Weighting Function W(n) for bounding box regression.

Args:
        max_num_bins (int): Max number of the discrete bins.
        up (Tensor): Controls upper bounds of the sequence,
                    where maximum offset is ±up * H / W.
        reg_scale (float): Controls the curvature of the Weighting Function.
                        Larger values result in flatter weights near the central axis W(max_num_bins/2)=0
                        and steeper weights at both ends.

Returns:
        Tensor: Sequence of Weighting Function.


def apply(self, predictions: Iterable[AnomalyPrediction]) -> AnomalyPrediction:
    result_dict: dict[str, Any] = {}
    _AggModelIdMixin.add_model_id(self, result_dict)
    _SourcePredictionMixin.add_source_predictions(self, result_dict, predictions)
    scores = [prediction.score for prediction in predictions if prediction.score is not None and (not math.isnan(prediction.score))]
    if len(scores) > 0:
        result_dict['score'] = self._agg(scores)
    elif all(map(lambda x: x.score is None, predictions)):
        result_dict['score'] = None
    else:
        result_dict['score'] = float('NaN')
    return AnomalyPrediction(**result_dict)

Applies the score aggregation function to a list of predictions.

Args:
        predictions (Iterable[AnomalyPrediction]): A collection of
        `AnomalyPrediction` objects to be aggregated.

Returns:
        AnomalyPrediction: A single `AnomalyPrediction` object with the
        aggregated score. The aggregated score is determined as follows:

        - If there are any non-missing and non-error scores, the `agg_func` is
        applied to aggregate them.
        - If all scores are error scores (`None`), the aggregated score is also
        `None`.
        - If there are a mix of missing (`NaN`) and error scores (`None`), the
        aggregated score is `NaN`.


def _variable_product_items(variableitems: Iterable[tuple[str, cfg.Variable]], complexity_limit: 'ComplexityLimit') -> Generator[list[tuple[str, cfg.Binding]], None, None]:
    variableitems_iter = iter(variableitems)
    try:
        headkey, headvar = next(variableitems_iter)
    except StopIteration:
        yield []
    else:
        for tail in _variable_product_items(variableitems_iter, complexity_limit):
            for headvalue in headvar.bindings:
                complexity_limit.inc()
                yield ([(headkey, headvalue)] + tail)

Take the Cartesian product of a list of (key, value) tuples.

    See variable_product_dict below.

Args:
        variableitems: A dict mapping object to cfg.Variable.
        complexity_limit: A counter that tracks how many combinations we've yielded
        and aborts if we go over the limit.

Yields:
        A sequence of [(key, cfg.Binding), ...] lists.


def __init__(self, indent, last_space):
    self.indent = indent
    self.last_space = last_space
    self.closing_scope_indent = 0
    self.split_before_closing_bracket = False
    self.num_line_splits = 0

Maintains the state of the bracket enclosures.

    A stack of _ParenState objects are kept so that we know how to indent relative
    to the brackets.

Attributes:
        indent: The column position to which a specified parenthesis level needs to
        be indented.
        last_space: The column position of the last space on each level.
        closing_scope_indent: The column position of the closing indentation.
        split_before_closing_bracket: Whether a newline needs to be inserted before
        the closing bracket. We only want to insert a newline before the closing
        bracket if there also was a newline after the beginning left bracket.
        num_line_splits: Number of line splits this _ParenState contains already.
        Each subsequent line split gets an increasing penalty.


def row_splits(self):
    return self._row_partition.row_splits()

The row-split indices for this ragged tensor's `values`.

    `rt.row_splits` specifies where the values for each row begin and end in
    `rt.values`.  In particular, the values for row `rt[i]` are stored in
    the slice `rt.values[rt.row_splits[i]:rt.row_splits[i+1]]`.

Returns:
        A 1-D integer `Tensor` with shape `[self.nrows+1]`.
        The returned tensor is non-empty, and is sorted in ascending order.
        `self.row_splits[0]` is zero, and `self.row_splits[-1]` is equal to
        `self.values.shape[0]`.

    #### Example:

    >>> rt = tf.ragged.constant([[3, 1, 4, 1], [], [5, 9, 2], [6], []])
    >>> print(rt.row_splits)  # indices of row splits in rt.values
    tf.Tensor([0 4 4 7 8 8], shape=(6,), dtype=int64)


def DefaultNodeVisit(self, node):
    for child in node.children:
        self.Visit(child)

Default visitor for Node: visits the node's children depth-first.

    This method is invoked when no specific visitor for the node is defined.

Args:
        node: the node to visit


def copy_clean(node, preserve_annos=None):
    return CleanCopier(preserve_annos).copy(node)

Creates a deep copy of an AST.

    The copy will not include fields that are prefixed by '__', with the
    exception of user-specified annotations.

Args:
        node: ast.AST
        preserve_annos: Optional[Set[Hashable]], annotation keys to include in the
        copy

Returns:
        ast.AST


def add(self, profile_datum):
    self.total_op_time += profile_datum.op_time
    self.total_exec_time += profile_datum.exec_time
    device_and_node = '%s:%s' % (profile_datum.device_name, profile_datum.node_exec_stats.node_name)
    device_and_node = '%s:%s' % (profile_datum.device_name, profile_datum.node_exec_stats.node_name)
    if device_and_node in self._node_to_exec_count:
        self._node_to_exec_count[device_and_node] += 1
    else:
        self._node_to_exec_count[device_and_node] = 1

Accumulate a new instance of ProfileDatum.

Args:
        profile_datum: (`ProfileDatum`) an instance of `ProfileDatum` to
        accumulate to this object.


def GetNetgroupMap(self, since=None):
    return NetgroupUpdateGetter().GetUpdates(self, self.conf['netgroup_url'], since)

Return the netgroup map from this source.

Args:
        since: Get data only changed since this timestamp (inclusive) or None
        for all data.

Returns:
        instance of netgroup.NetgroupMap


def terminate(self, task_type, task_id):
    with self._process_lock:
        p = self._processes.get((task_type, task_id), None)
        if p is None:
            raise ValueError('{}-{} does not exist'.format(task_type, task_id))
        self._terminated.add((task_type, task_id))
        self._parent_to_sub_queue.put('terminate {} {}'.format(task_type, task_id))
        p.join()

Terminates the process with `task_type` and `task_id`.

    If auto_retart=True, the terminated task will be restarted unless the chief
    has already exited with zero exit code.

Args:
        task_type: the task type.
        task_id: the task id.


def _parse_unknown_block_line(self, instrumentation_block, line):
    if line.startswith(_InstrumentationStructurePrefixes.STATUS):
        return self._parse_method_block_line(self._transition_instrumentation_block(instrumentation_block, new_state=_InstrumentationBlockStates.METHOD), line)
    elif line.startswith(_InstrumentationStructurePrefixes.RESULT) or _InstrumentationStructurePrefixes.FAILED in line:
        return self._parse_result_block_line(self._transition_instrumentation_block(instrumentation_block, new_state=_InstrumentationBlockStates.RESULT), line)
    else:
        instrumentation_block.add_value(line)
        return instrumentation_block

Parses a line from the instrumentation output from the UNKNOWN
    parser state.

Args:
        instrumentation_block: _InstrumentationBlock, the current
        instrumenation block, where the correct categorization it noti
        yet known.
        line: string, the raw instrumenation output line to be used to
        deteremine the correct categorization.

Returns:
        The next instrumentation block to continue parsing with. Usually,
        this is the same instrumentation block but with the state
        transitioned appropriately.


def __call__(self, images: ImageInput=None, text: Optional[Union[List[str], List[List[str]]]]=None, audio=None, videos=None, **kwargs: Unpack[OmDetTurboProcessorKwargs]) -> BatchFeature:
    if images is None or text is None:
        raise ValueError('You have to specify both `images` and `text`')
    output_kwargs = self._merge_kwargs(OmDetTurboProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs)
    if isinstance(text, str):
        text = text.strip(' ').split(',')
    if not (len(text) and isinstance(text[0], (list, tuple))):
        text = [text]
    task = output_kwargs['text_kwargs'].pop('task', None)
    if task is None:
        task = ['Detect {}.'.format(', '.join(text_single)) for text_single in text]
    elif not isinstance(task, (list, tuple)):
        task = [task]
    encoding_image_processor = self.image_processor(images, **output_kwargs['images_kwargs'])
    tasks_encoding = self.tokenizer(text=task, **output_kwargs['text_kwargs'])
    classes = text
    classes_structure = torch.tensor([len(class_single) for class_single in classes], dtype=torch.long)
    classes_flattened = [class_single for class_batch in classes for class_single in class_batch]
    classes_encoding = self.tokenizer(text=classes_flattened, **output_kwargs['text_kwargs'])
    encoding = BatchFeature()
    encoding.update({f'tasks_{key}': value for key, value in tasks_encoding.items()})
    encoding.update({f'classes_{key}': value for key, value in classes_encoding.items()})
    encoding.update({'classes_structure': classes_structure})
    encoding.update(encoding_image_processor)
    return encoding

This method uses [*DetrImageProcessor.__call__] method to prepare image(s) for the model, and
    [CLIPTokenizerFast.__call__] to prepare text for the model.

    Please refer to the docstring of the above two methods for more information.

Args:
        images (`ImageInput`):
        Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255.
        text (`Union[str, List[str], List[List[str]]]`):
            The classes used to limit the scope of the open vocabulary detection. Expects a list of strings or a list
            of list of strings. Batched classes can be of different lengths.
            Examples: ["cat", "dog", "bird"], [["cat", "dog", "bird"], ["hat", "person"], ["car"]]
    Kwargs:
        task (`Union[str, List[str], TextInput, PreTokenizedInput]`):
            The grounded text used to guide open vocabulary detection. Expects a single string or a list of strings.
            Examples: "Detect a cat, a dog, and a bird.",[ "Detect everything.", "Detect trees and flowers."]
            When not provided, the default task is "Detect [class1], [class2], [class3]" etc.
        ...


def deserialize(config, custom_objects=None):
    return deserialize_keras_object(config, module_objects=globals(), custom_objects=custom_objects, printable_module_name='metric function')

Deserializes a serialized metric class/function instance.

Args:
        config: Metric configuration.
        custom_objects: Optional dictionary mapping names (strings) to custom
        objects (classes and functions) to be considered during deserialization.

Returns:
        A Keras `Metric` instance or a metric function.


def apply_indexed_slices_grad(self, grad, local_step=0, name=None):
    return self.apply_grad(grad_indices=grad.indices, grad_values=grad.values, grad_shape=grad.dense_shape, local_step=local_step, name=name)

Attempts to apply a gradient to the accumulator.

    The attempt is silently dropped if the gradient is stale, i.e., `local_step`
    is less than the accumulator's global time step.

Args:
        grad: The gradient `IndexedSlices` to be applied.
        local_step: Time step at which the gradient was computed.
        name: Optional name for the operation.

Returns:
        The operation that (conditionally) applies a gradient to the accumulator.

Raises:
        InvalidArgumentError: If grad is of the wrong shape


def _set_variable_or_list_initializer(variable_or_list, ckpt_file, tensor_name):
    if isinstance(variable_or_list, (list, tuple)):
        slice_name = None
        for v in variable_or_list:
            slice_info = v._save_slice_info
            if slice_name is None:
                slice_name = slice_info.full_name
            elif slice_name != slice_info.full_name:
                raise ValueError('Slices must all be from the same tensor: %s != %s' % (slice_name, slice_info.full_name))
            _set_checkpoint_initializer(v, ckpt_file, tensor_name, slice_info.spec)
    else:
        _set_checkpoint_initializer(variable_or_list, ckpt_file, tensor_name, '')

Overrides initialization op of given variable or list of variables.

    Calls `_set_checkpoint_initializer` for each variable in the given list of
    variables.

Args:
        variable_or_list: `tf.Variable` object or a list of `tf.Variable` objects.
        ckpt_file: string, full path of the checkpoint.
        tensor_name: Name of the tensor to load from the checkpoint.

Raises:
        ValueError: if all objects in `variable_or_list` are not partitions of the
        same large variable.


def auto_call_functors(enabled: bool=True) -> ContextManager[None]:
    return thread_local.thread_local_value_scope(_TLS_AUTO_CALL_FUNCTORS, enabled, False)

Returns a context manager to enable or disable auto call for functors.

    `auto_call_functors` is thread-safe and can be nested. For example::

        @pg.symbolize
        def foo(x, y):
        return x + y

        with pg.auto_call_functors(True):
        a = foo(1, 2)
        assert a == 3
        with pg.auto_call_functors(False):
        b = foo(1, 2)
        assert isinstance(b, foo)

Args:
        enabled: If True, enable auto call for functors.
        Otherwise, auto call will be disabled.

Returns:
        A context manager for enabling/disabling auto call for functors.


def call(self, input_features: TFModelInputType | None=None, decoder_input_ids: np.ndarray | tf.Tensor | None=None, decoder_attention_mask: np.ndarray | tf.Tensor | None=None, decoder_position_ids: np.ndarray | tf.Tensor | None=None, head_mask: np.ndarray | tf.Tensor | None=None, decoder_head_mask: np.ndarray | tf.Tensor | None=None, cross_attn_head_mask: np.ndarray | tf.Tensor | None=None, encoder_outputs: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]]=None, past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]]=None, decoder_inputs_embeds: Optional[Tuple[Union[np.ndarray, tf.Tensor]]]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, training: bool=False) -> Union[Tuple[tf.Tensor], TFSeq2SeqModelOutput]:
    outputs = self.model(input_features=input_features, decoder_input_ids=decoder_input_ids, decoder_attention_mask=decoder_attention_mask, decoder_position_ids=decoder_position_ids, head_mask=head_mask, decoder_head_mask=decoder_head_mask, cross_attn_head_mask=cross_attn_head_mask, encoder_outputs=encoder_outputs, past_key_values=past_key_values, decoder_inputs_embeds=decoder_inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, training=training)
    return outputs

Returns:

Example:
    ```python
    >>> import tensorflow as tf
    >>> from transformers import TFWhisperModel, AutoFeatureExtractor
    >>> from datasets import load_dataset

    >>> model = TFWhisperModel.from_pretrained("openai/whisper-base")
    >>> feature_extractor = AutoFeatureExtractor.from_pretrained("openai/whisper-base")
    >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
    >>> inputs = feature_extractor(ds[0]["audio"]["array"], return_tensors="tf")
    >>> input_features = inputs.input_features
    >>> decoder_input_ids = tf.convert_to_tensor([[1, 1]]) * model.config.decoder_start_token_id
    >>> last_hidden_state = model(input_features, decoder_input_ids=decoder_input_ids).last_hidden_state
    >>> list(last_hidden_state.shape)
    [1, 2, 512]
    ```


def resize(self, image: np.ndarray, size: Union[Dict[str, int], int], resample: PILImageResampling=PILImageResampling.BICUBIC, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None, **kwargs) -> np.ndarray:
    if input_data_format is None:
        input_data_format = infer_channel_dimension_format(image)
    height, width = get_image_size(image, input_data_format)
    max_size = max(height, width)
    size = get_size_dict(size, default_to_square=True)
    if size['height'] != size['width']:
        raise ValueError(f'Output height and width must be the same. Got height={size['height']} and width={size['width']}')
    size = size['height']
    delta = size / max_size
    output_size_nonpadded = [max(int(height * delta), self.min_size), max(int(width * delta), self.min_size)]
    image = resize(image, size=output_size_nonpadded, resample=resample, data_format=data_format, input_data_format=input_data_format, **kwargs)
    image = self.pad_to_square(image=image, background_color=self.background_color, input_data_format=input_data_format)
    return image

Resize an image to dynamically calculated size.

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

Returns:
        `np.ndarray`: The resized image.


def _get_split_key(client_keys, num_splits):
    if not client_keys or len(client_keys) < num_splits - 1:
        return client_keys
    num_keys_per_split = max(1.0, float(len(client_keys)) / (num_splits - 1))
    split_client_keys = []
    for i in range(1, num_splits):
        split_index = int(round(i * num_keys_per_split) - 1)
        split_client_keys.append(client_keys[split_index])
    return split_client_keys

Given a list of keys and a number of splits find the keys to split on.

Args:
        client_keys: the list of keys.
        num_splits: the number of splits.

Returns:
        A list of keys to split on.


def _escape_identifier(identifier_value: str) -> str:
    identifier_value = identifier_value.replace('-', '_')
    if identifier_value.upper() in _sql_data_types.STANDARD_SQL_KEYWORDS:
        return f'`{identifier_value}`'
    return identifier_value

Returns the value in a valid column name format.

Args:
        identifier_value: Raw identifier string.

Returns:
        If the identifier is a keyword, then it will be surrounded by backticks. If
        the identifier contains a hyphen, the hyphen will be replaced by an
        underscore.


def custom_apply(self, path: utils.KeyPath, value_spec: pg_typing.ValueSpec, allow_partial: bool, child_transform: Optional[Callable[[utils.KeyPath, pg_typing.Field, Any], Any]]=None) -> Tuple[bool, Any]:
    del path, value_spec, allow_partial, child_transform
    return (False, self)

Custom apply on a value based on its original value spec.

    This implements ``pg.pg_typing.CustomTyping``, allowing a pure symbolic
    value to be assigned to any field. To customize this behavior, override
    this method in subclasses.

Args:
        path: KeyPath of current object under its object tree.
        value_spec: Original value spec for this field.
        allow_partial: Whether allow partial object to be created.
        child_transform: Function to transform child node values into their final
        values. Transform function is called on leaf nodes first, then on their
        parents, recursively.

Returns:
        A tuple (proceed_with_standard_apply, value_to_proceed).
        If proceed_with_standard_apply is set to False, value_to_proceed
        will be used as final value.

Raises:
        Error when the value is not compatible with the value spec.


def resource_schema(self, resource: str) -> dict:
    entry = self.api_document['schemas'][resource]['properties']
    return self.to_schema(entry)

Return BigQuery schema for a Discovery API resource.

Args:
        resource: the name of the Google API resource

Returns:
        A dictionary representation of the resource.


def _findSourceLine(self, annotated_source, line_number):
    index = None
    for i, line in enumerate(annotated_source.lines):
        if line.startswith('L%d ' % line_number):
            index = i
            break
    return index

Find line of given line number in annotated source.

Args:
        annotated_source: (debugger_cli_common.RichTextLines) the annotated source
        line_number: (int) 1-based line number

Returns:
        (int) If line_number is found, 0-based line index in
        annotated_source.lines. Otherwise, None.


def eye(size, dtype=None, name=None):
    if dtype is None:
        dtype = floatx()
    tf_dtype = dtypes_module.as_dtype(dtype)
    return variable(linalg_ops.eye(size, dtype=tf_dtype), dtype, name)

Instantiate an identity matrix and returns it.

Args:
        size: Integer, number of rows/columns.
        dtype: String, data type of returned Keras variable.
        name: String, name of returned Keras variable.

Returns:
        A Keras variable, an identity matrix.

Example:
    >>> kvar = tf.keras.backend.eye(3)
    >>> tf.keras.backend.eval(kvar)
    array([[1.,  0.,  0.],
        [0.,  1.,  0.],
        [0.,  0.,  1.]], dtype=float32)


def run_bq_pipeline(argv=None):
    parser = argparse.ArgumentParser()
    parser.add_argument('--query', required=True, help='Query to process for the table.')
    parser.add_argument('--output', required=True, help='Output BQ table to write results to.')
    parser.add_argument('--output_schema', dest='output_schema', required=True, help='Schema for output BQ table.')
    parser.add_argument('--use_standard_sql', action='store_true', dest='use_standard_sql', help='Output BQ table to write results to.')
    parser.add_argument('--kms_key', default=None, help='Use this Cloud KMS key with BigQuery.')
    parser.add_argument('--native', default=False, action='store_true', help='Use NativeSources and Sinks.')
    parser.add_argument('--use_json_exports', default=False, action='store_true', help='Use JSON as the file format for exports.')
    known_args, pipeline_args = parser.parse_known_args(argv)
    table_schema = parse_table_schema_from_json(known_args.output_schema)
    kms_key = known_args.kms_key
    options = PipelineOptions(pipeline_args)
    p = TestPipeline(options=options)
    if known_args.native:
        data = p | 'read' >> beam.io.Read(beam.io.BigQuerySource(query=known_args.query, use_standard_sql=known_args.use_standard_sql, kms_key=kms_key))
    else:
        data = p | 'read' >> beam.io.gcp.bigquery.ReadFromBigQuery(query=known_args.query, project=options.view_as(GoogleCloudOptions).project, use_standard_sql=known_args.use_standard_sql, use_json_exports=known_args.use_json_exports, kms_key=kms_key)
    temp_file_format = 'NEWLINE_DELIMITED_JSON' if known_args.use_json_exports else 'AVRO'
    _ = data | 'write' >> beam.io.WriteToBigQuery(known_args.output, schema=table_schema, create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=beam.io.BigQueryDisposition.WRITE_EMPTY, temp_file_format=temp_file_format, kms_key=kms_key)
    result = p.run()
    result.wait_until_finish()

Run the sample BigQuery pipeline.

Args:
        argv: Arguments to the run function.


def with_max_depth(self, max_depth):
    self._options['max_depth'] = max_depth
    return self

Set the maximum depth of display.

    The depth depends on profiling view. For 'scope' view, it's the
    depth of name scope hierarchy (tree), for 'op' view, it's the number
    of operation types (list), etc.

Args:
        max_depth: Maximum depth of the data structure to display.

Returns:
        self