def _filter_valid_filepaths(self, df, x_col):
    filepaths = df[x_col].map(lambda fname: os.path.join(self.directory, fname))
    mask = filepaths.apply(validate_filename, args=(self.white_list_formats,))
    n_invalid = (~mask).sum()
    if n_invalid:
        warnings.warn('Found {} invalid image filename(s) in x_col="{}". These filename(s) will be ignored.'.format(n_invalid, x_col))
    return df[mask]

Keep only dataframe rows with valid filenames.

Args:
        df: Pandas dataframe containing filenames in a column
        x_col: string, column in `df` that contains the filenames or
            filepaths

Returns:
        absolute paths to image files


def __init__(self, x: int, y: int=1, **kwargs):
    self.z = x + y + sum(kwargs.values())

Class A.

Args:
        x: The first integer.
        y: The second integer.
        **kwargs: Other arguments.


def __init__(self, config, n_bins=16, mlp_dim=256, min_depth=0.001, max_depth=10):
    super().__init__()
    self.in_features = config.bottleneck_features
    self.bin_centers_type = config.bin_centers_type
    self.min_depth = min_depth
    self.max_depth = max_depth
    self.conv1 = nn.Conv2d(self.in_features, mlp_dim, 1, 1, 0)
    self.act1 = nn.ReLU(inplace=True)
    self.conv2 = nn.Conv2d(mlp_dim, n_bins, 1, 1, 0)
    self.act2 = nn.ReLU(inplace=True) if self.bin_centers_type == 'normed' else nn.Softplus()

Bin center regressor network.

    Can be "normed" or "unnormed". If "normed", bin centers are bounded on the (min_depth, max_depth) interval.

Args:
        config (`int`):
            Model configuration.
        n_bins (`int`, *optional*, defaults to 16):
            Number of bin centers.
        mlp_dim (`int`, *optional*, defaults to 256):
            Hidden dimension.
        min_depth (`float`, *optional*, defaults to 1e-3):
            Min depth value.
        max_depth (`float`, *optional*, defaults to 10):
            Max depth value.


def run(argv=None):
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', dest='input', default='gs://dataflow-samples/wikipedia_edits/*.json', help='Input specified as a GCS path containing a BigQuery table exported as json.')
    parser.add_argument('--output', required=True, help='Output file to write results to.')
    parser.add_argument('--sampling_threshold', type=float, default=0.1, help='Fraction of entries used for session tracking')
    known_args, pipeline_args = parser.parse_known_args(argv)
    pipeline_options = PipelineOptions(pipeline_args)
    pipeline_options.view_as(SetupOptions).save_main_session = True
    with beam.Pipeline(options=pipeline_options) as p:
        p | ReadFromText(known_args.input) | ComputeTopSessions(known_args.sampling_threshold) | WriteToText(known_args.output)

Runs the Wikipedia top edits pipeline.

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


def cumsum(x, axis=None, dtype=None):
    return Cumsum(axis=axis, dtype=dtype)(x)

Returns the cumulative sum of elements along a given axis.

Args:
        x: Input tensor.
        axis: Axis along which the cumulative sum is computed.
            By default the input is flattened.
        dtype: dtype of returned tensor. Defaults to x.dtype.

Returns:
        Output tensor.


def remat(f):

    def wrapped(*args, **kwargs):
        return torch.utils.checkpoint.checkpoint(f, *args, use_reentrant=False)
    return wrapped

Implementation of rematerialization.

Args:
        f: The function or operation to rematerialize.

Returns:
        A function wrapping f that defines a custom gradient, which
        recomputes f on the backwards pass of a gradient call.


def GetConsoleAttr(encoding=None, reset=False):
    attr = ConsoleAttr._CONSOLE_ATTR_STATE
    if not reset:
        if not attr:
            reset = True
        elif encoding and encoding != attr.GetEncoding():
            reset = True
    if reset:
        attr = ConsoleAttr(encoding=encoding)
        ConsoleAttr._CONSOLE_ATTR_STATE = attr
    return attr

Gets the console attribute state.

    If this is the first call or reset is True or encoding is not None and does
    not match the current encoding or out is not None and does not match the
    current out then the state is (re)initialized. Otherwise the current state
    is returned.

    This call associates the out file stream with the console. All console related
    output should go to the same stream.

Args:
        encoding: Encoding override.
        ascii -- ASCII. This is the default.
        utf8 -- UTF-8 unicode.
        win -- Windows code page 437.
        reset: Force re-initialization if True.

Returns:
        The global ConsoleAttr state object.


def _einsum_v1_reduction(t0, t0_axis_labels, t1, t1_axis_labels, axes_to_sum):
    if len(t0_axis_labels) != len(t0.shape):
        raise ValueError(f'Tensor `t0` of rank {len(t0.shape)} does not match einsum reduction of length {len(t0_axis_labels)}.')
    if len(t1_axis_labels) != len(t1.shape):
        raise ValueError(f'Tensor `t1` of rank {len(t1.shape)} does not match einsum reduction of length {len(t1_axis_labels)}')
    assert all((a in t0_axis_labels and a in t1_axis_labels for a in axes_to_sum))
    preserved_axes = (set(t0_axis_labels) & set(t1_axis_labels)) - axes_to_sum
    broadcast_axes = {}
    for i, sym_list in enumerate([t0_axis_labels, t1_axis_labels]):
        broadcast_axes[i] = set(sym_list) - preserved_axes - axes_to_sum

    def sort_key(input_index, a):
        if a in preserved_axes:
            return (-1, a)
        elif input_index == 0 and a in broadcast_axes[0] or (input_index == 1 and a in axes_to_sum):
            return (0, a)
        else:
            return (1, a)
    axis_labels = [t0_axis_labels, t1_axis_labels]
    sorted_axes = [sorted(sym_list, key=lambda a: sort_key(i, a)) for i, sym_list in enumerate(axis_labels)]
    inputs = [t0, t1]
    for i, axes_str in enumerate(axis_labels):
        perm = [axes_str.find(a) for a in sorted_axes[i]]
        inputs[i] = _transpose_if_necessary(inputs[i], perm)
    t0, t1 = inputs
    if not axes_to_sum:
        for _ in broadcast_axes[1]:
            t0 = array_ops.expand_dims(t0, -1)
        for _ in broadcast_axes[0]:
            t1 = array_ops.expand_dims(t1, len(preserved_axes))
        product = math_ops.multiply(t0, t1)
        product_axes = sorted_axes[0] + sorted_axes[1][len(preserved_axes):]
        return (product, ''.join(product_axes))
    else:
        t0_shape = _get_shape(t0)
        num_broadcast_elements_t0 = _total_size(t0_shape[len(preserved_axes):-len(axes_to_sum)])
        num_summed_elements = _total_size(t0_shape[-len(axes_to_sum):])
        new_shape = t0_shape[:len(preserved_axes)] + [num_broadcast_elements_t0, num_summed_elements]
        t0 = _reshape_if_necessary(t0, new_shape)
        t1_shape = _get_shape(t1)
        num_broadcast_elements_t1 = _total_size(t1_shape[len(preserved_axes) + len(axes_to_sum):])
        new_shape = t1_shape[:len(preserved_axes)] + [num_summed_elements, num_broadcast_elements_t1]
        t1 = _reshape_if_necessary(t1, new_shape)
        product = math_ops.matmul(t0, t1)
        uncompacted_shape = t0_shape[:len(preserved_axes) + len(broadcast_axes[0])] + t1_shape[len(t1_shape) - len(broadcast_axes[1]):]
        product = _reshape_if_necessary(product, uncompacted_shape)
        product_axes = sorted_axes[0][:len(preserved_axes) + len(broadcast_axes[0])] + sorted_axes[1][len(sorted_axes[1]) - len(broadcast_axes[1]):]
        return (product, ''.join(product_axes))

Helper for einsum() that computes the result of a two-argument einsum().

Args:
        t0: a `Tensor`
        t0_axis_labels: a string of axis labels.  This string's length must equal
        the rank of t0.
        t1: a `Tensor`
        t1_axis_labels: a string to axis labels.  This string's length must equal
        the rank of t1.
        axes_to_sum: set of labels of axes to be summed over

Returns:
        A `Tensor` whose elements are obtained by summing, over all axes in
        `axes_to_sum`, the corresponding elements of `t0` and `t1`.

        For example, if t0_axis_labels == 'abijk', t1_axis_labels == 'acjkl', and
        axes_to_sum == {j,k}, this will return a tensor x where

        out[a,b,c,i,l] = sum_j sum_k t0[a,b,i,j,k] * t1[a,c,j,k,l]

Raises:
        ValueError: if the rank of `t0` does not match the length of
        `t0_axis_labels`, or that of `t1` does not match the length of
        `t1_axis_labels`.


def chain(processor_list: Sequence[Processor | PartProcessor]) -> Processor:
    if not processor_list:
        raise ValueError('processor_list is empty')
    chain_processor = processor_list[0]
    for p in processor_list[1:]:
        chain_processor = chain_processor + p
    if isinstance(chain_processor, PartProcessor):
        chain_processor = chain_processor.to_processor()
    return chain_processor

Chain a sequence of processors.

Args:
        processor_list: list of part processors or generic processors.

Returns:
        A processor consisting of the chain of all the processors in the list. The
        execution is sequential from the first processor to the last but parts are
        processed concurrently overall.


def __init__(self, tokenizer=None, **kwargs):
    if tokenizer is None:
        raise ValueError('You need to specify a `tokenizer`.')
    super().__init__(tokenizer)

Constructs a Bros processor which wraps a BERT tokenizer.

    [`BrosProcessor`] offers all the functionalities of [`BertTokenizerFast`]. See the docstring of
    [`~BrosProcessor.__call__`] and [`~BrosProcessor.decode`] for more information.

Args:
        tokenizer (`BertTokenizerFast`, *optional*):
            An instance of ['BertTokenizerFast`]. The tokenizer is a required input.


def _skip_parameter_matching(self) -> bool:
    if self.signature.type_params:
        return False
    if self.ctx.options.analyze_annotated:
        return False
    return self.signature.has_return_annotation or self.full_name == '__init__'

Check whether we should skip parameter matching.

    This is use to skip parameter matching for function calls in the context of
    inference (pyi generation). This is to optimize the case where we don't
    need to match parameters in cases which the function has explicit type
    annotations, meaning that we don't need to infer the type.

Returns:
        True if we should skip parameter matching.


def delete(paths):
    if isinstance(paths, str):
        raise BeamIOError('Delete passed string argument instead of list: %s' % paths)
    if len(paths) == 0:
        return
    filesystem = FileSystems.get_filesystem(paths[0])
    return filesystem.delete(paths)

Deletes files or directories at the provided paths.
    Directories will be deleted recursively.

Args:
        paths: list of paths that give the file objects to be deleted

Raises:
        ``BeamIOError``: if any of the delete operations fail


def max_pooling3d(inputs, pool_size, strides, padding='valid', data_format='channels_last', name=None):
    warnings.warn('`tf.layers.max_pooling3d` is deprecated and will be removed in a future version. Please use `tf.keras.layers.MaxPooling3D` instead.')
    layer = MaxPooling3D(pool_size=pool_size, strides=strides, padding=padding, data_format=data_format, name=name)
    return layer.apply(inputs)

Max pooling layer for 3D inputs (e.g.

    volumes).

Args:
        inputs: The tensor over which to pool. Must have rank 5.
        pool_size: An integer or tuple/list of 3 integers: (pool_depth, pool_height,
        pool_width) specifying the size of the pooling window. Can be a single
        integer to specify the same value for all spatial dimensions.
        strides: An integer or tuple/list of 3 integers, specifying the strides of
        the pooling operation. Can be a single integer to specify the same value
        for all spatial dimensions.
        padding: A string. The padding method, either 'valid' or 'same'.
        Case-insensitive.
        data_format: A string. The ordering of the dimensions in the inputs.
        `channels_last` (default) and `channels_first` are supported.
        `channels_last` corresponds to inputs with shape `(batch, depth, height,
        width, channels)` while `channels_first` corresponds to inputs with shape
        `(batch, channels, depth, height, width)`.
        name: A string, the name of the layer.

Returns:
        Output tensor.

Raises:
        ValueError: if eager execution is enabled.


def compute_column_logits(sequence_output, column_output_weights, column_output_bias, cell_index, cell_mask, allow_empty_column_selection):
    token_logits = torch.einsum('bsj,j->bs', sequence_output, column_output_weights) + column_output_bias
    cell_logits, cell_logits_index = reduce_mean(token_logits, cell_index)
    column_index = cell_index.project_inner(cell_logits_index)
    column_logits, out_index = reduce_sum(cell_logits * cell_mask, column_index)
    cell_count, _ = reduce_sum(cell_mask, column_index)
    column_logits /= cell_count + EPSILON_ZERO_DIVISION
    is_padding = torch.logical_and(cell_count < 0.5, ~torch.eq(out_index.indices, 0))
    column_logits += CLOSE_ENOUGH_TO_LOG_ZERO * torch.as_tensor(is_padding, dtype=torch.float32, device=is_padding.device)
    if not allow_empty_column_selection:
        column_logits += CLOSE_ENOUGH_TO_LOG_ZERO * torch.as_tensor(torch.eq(out_index.indices, 0), dtype=torch.float32, device=out_index.indices.device)
    return column_logits

Computes the column logits.

Args:
        sequence_output (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
            Also known as last_hidden_state. Sequence of hidden-states at the output of the last layer of the model.
        column_output_weights (`torch.FloatTensor` of shape `(hidden_size)`):
            Weights of the linear layer for column selection.
        column_output_bias (`torch.FloatTensor` of shape `()`):
            Bias of the linear layer for column selection.
        cell_index (`ProductIndexMap`):
            Index that groups tokens into cells.
        cell_mask (`torch.FloatTensor` of shape `(batch_size, max_num_rows * max_num_cols)`):
            Mask for cells that exist in the table (i.e. that are not padding).
        allow_empty_column_selection (`bool`):
            Whether to allow not to select any column

Returns:
        column_logits (`torch.FloatTensor`of shape `(batch_size, max_num_cols)`): Tensor containing the column logits
        for every example in the batch.


def _get_dequantized_hist_mids_after_quantize(self, quant_min: float, quant_max: float) -> np.ndarray:
    maxbound = 2 ** self._num_bits - 1
    minbound = 0
    scale = (quant_max - quant_min) / maxbound
    zero_point = -quant_min / scale
    if abs(zero_point) > 9000000000.0:
        zero_point = 9000000000.0
    if abs(scale) < 1e-09:
        scale = 1e-09
    zero_point = round(zero_point)
    quantized_hist_mids = np.clip(np.round(self._hist_mids / scale) + zero_point, minbound, maxbound)
    dequantized_hist_mids = scale * (quantized_hist_mids - zero_point)
    return dequantized_hist_mids

Quantizes and dequantizes hist_mids using quant_min and quant_max.

    Quantization converts the range of numbers from [quant_min, quant_max] to
    [0, 2^num_bits - 1]. Values less than quant_min are converted to 0, and
    values greater than quant_max are converted to 2^num_bits - 1.

    The histogram represents the distribution of the data, and our goal is to
    find the quant_min and quant_max that best describe this distribution. To do
    this, we quantize hist_mids using quant_min and quant_max and dequantize
    them again. Then the difference between hist_mids and dequantized hist_mids
    equates to quantization error when using quant_min and quant_max.

Args:
        quant_min: The minimum real value that can be represented by a quantized
        value.
        quant_max: The maximum real value that can be represented by a quantized
        value.

Returns:
        dequantized hist_mids after quantizing by quant_min and quant_max


def slice_arrays(arrays, start=None, stop=None):
    if arrays is None:
        return [None]
    if isinstance(start, list) and stop is not None:
        raise ValueError('The stop argument has to be None if the value of start is a list.')
    elif isinstance(arrays, list):
        if hasattr(start, '__len__'):
            if hasattr(start, 'shape'):
                start = start.tolist()
            return [None if x is None else x[start] for x in arrays]
        return [None if x is None else None if not hasattr(x, '__getitem__') else x[start:stop] for x in arrays]
    else:
        if hasattr(start, '__len__'):
            if hasattr(start, 'shape'):
                start = start.tolist()
            return arrays[start]
        if hasattr(start, '__getitem__'):
            return arrays[start:stop]
        return [None]

Slice an array or list of arrays.

    This takes an array-like, or a list of
    array-likes, and outputs:
        - arrays[start:stop] if `arrays` is an array-like
        - [x[start:stop] for x in arrays] if `arrays` is a list

    Can also work on list/array of indices: `slice_arrays(x, indices)`

Args:
        arrays: Single array or list of arrays.
        start: can be an integer index (start index) or a list/array of indices
        stop: integer (stop index); should be None if `start` was a list.

Returns:
        A slice of the array(s).

Raises:
        ValueError: If the value of start is a list and stop is not None.


def rot90(array, k=1, axes=(0, 1)):
    if any_symbolic_tensors((array,)):
        return Rot90(k=k, axes=axes).symbolic_call(array)
    return backend.numpy.rot90(array, k=k, axes=axes)

Rotate an array by 90 degrees in the plane specified by axes.

    This function rotates an array counterclockwise
    by 90 degrees `k` times in the plane specified by `axes`.
    Supports arrays of two or more dimensions.

Args:
        array: Input array to rotate.
        k: Number of times the array is rotated by 90 degrees.
        axes: A tuple of two integers specifying the
            plane of rotation (defaults to `(0, 1)`).

Returns:
        Rotated array.

Example:
    >>> import numpy as np
    >>> from keras import ops
    >>> m = np.array([[1, 2], [3, 4]])
    >>> rotated = ops.rot90(m)
    >>> rotated
    array([[2, 4],
        [1, 3]])

    >>> m = np.arange(8).reshape((2, 2, 2))
    >>> rotated = ops.rot90(m, k=1, axes=(1, 2))
    >>> rotated
    array([[[1, 3],
            [0, 2]],
        [[5, 7],
            [4, 6]]])


def load_img(path, color_mode='rgb', target_size=None, interpolation='nearest', keep_aspect_ratio=False):
    if pil_image is None:
        raise ImportError('Could not import PIL.Image. The use of `load_img` requires PIL.')
    if isinstance(path, io.BytesIO):
        img = pil_image.open(path)
    elif isinstance(path, (pathlib.Path, bytes, str)):
        if isinstance(path, pathlib.Path):
            path = str(path.resolve())
        with open(path, 'rb') as f:
            img = pil_image.open(io.BytesIO(f.read()))
    else:
        raise TypeError(f'path should be path-like or io.BytesIO, not {type(path)}')
    if color_mode == 'grayscale':
        if img.mode not in ('L', 'I;16', 'I'):
            img = img.convert('L')
    elif color_mode == 'rgba':
        if img.mode != 'RGBA':
            img = img.convert('RGBA')
    elif color_mode == 'rgb':
        if img.mode != 'RGB':
            img = img.convert('RGB')
    else:
        raise ValueError('color_mode must be "grayscale", "rgb", or "rgba"')
    if target_size is not None:
        width_height_tuple = (target_size[1], target_size[0])
        if img.size != width_height_tuple:
            if interpolation not in PIL_INTERPOLATION_METHODS:
                raise ValueError('Invalid interpolation method {} specified. Supported methods are {}'.format(interpolation, ', '.join(PIL_INTERPOLATION_METHODS.keys())))
            resample = PIL_INTERPOLATION_METHODS[interpolation]
            if keep_aspect_ratio:
                width, height = img.size
                target_width, target_height = width_height_tuple
                crop_height = width * target_height // target_width
                crop_width = height * target_width // target_height
                crop_height = min(height, crop_height)
                crop_width = min(width, crop_width)
                crop_box_hstart = (height - crop_height) // 2
                crop_box_wstart = (width - crop_width) // 2
                crop_box_wend = crop_box_wstart + crop_width
                crop_box_hend = crop_box_hstart + crop_height
                crop_box = [crop_box_wstart, crop_box_hstart, crop_box_wend, crop_box_hend]
                img = img.resize(width_height_tuple, resample, box=crop_box)
            else:
                img = img.resize(width_height_tuple, resample)
    return img

Loads an image into PIL format.

Example:
    ```python
    image = keras.utils.load_img(image_path)
    input_arr = keras.utils.img_to_array(image)
    input_arr = np.array([input_arr])  # Convert single image to a batch.
    predictions = model.predict(input_arr)
    ```

Args:
        path: Path to image file.
        color_mode: One of `"grayscale"`, `"rgb"`, `"rgba"`. Default: `"rgb"`.
            The desired image format.
        target_size: Either `None` (default to original size) or tuple of ints
            `(img_height, img_width)`.
        interpolation: Interpolation method used to resample the image if the
            target size is different from that of the loaded image. Supported
            methods are `"nearest"`, `"bilinear"`, and `"bicubic"`.
            If PIL version 1.1.3 or newer is installed, `"lanczos"`
            is also supported. If PIL version 3.4.0 or newer is installed,
            `"box"` and `"hamming"` are also
            supported. By default, `"nearest"` is used.
        keep_aspect_ratio: Boolean, whether to resize images to a target
            size without aspect ratio distortion. The image is cropped in
            the center with target aspect ratio before resizing.

Returns:
        A PIL Image instance.


def get_rots(self) -> Rotation:
    return self._rots

Getter for the rotation.

Returns:
        The rotation object


def __init__(self, conf):
    super(S3FilesSource, self).__init__(conf)
    self._SetDefaults(conf)
    self.s3_client = None

Initialise the S3FilesSource object.

Args:
        conf: A dictionary of key/value pairs.

Raises:
        RuntimeError: object wasn't initialised with a dict


def delete_weight(self, object_name, weight_name):

    def delete_weight_fn(weights_dict, source_name, target_name=None):
        if weight_name not in weights_dict[source_name]:
            raise ValueError(f'Weight {weight_name} not found in object {object_name}. Weights found: {list(weights_dict[source_name].keys())}')
        weights_dict[source_name].pop(weight_name)
    self._edit_object(delete_weight_fn, object_name)

Removes a weight from an existing object.

Args:
        object_name: String, name or path of the
            object from which to remove the weight
            (e.g. `"dense_2"` or `"layers/dense_2"`).
        weight_name: String, name of the weight to
            delete (e.g. `"0"`).


def set_base_log(self, dqm_version_id: str, workflow_execution_id: str, table_metadata: TableMetadata, run_dt_utc: datetime) -> None:
    self._base_log = LogMessage(dqm_version_id=dqm_version_id, workflow_execution_id=workflow_execution_id, project_id=table_metadata.project_id, dataset_id=table_metadata.dataset_id, table_name=table_metadata.table_name, full_table_id=table_metadata.full_table_id, run_timestamp_utc=get_formatted_timestamp(run_dt_utc))

Set the common values for each log message.

Args:
        * dqm_version_id: DQM version ID
        * workflow_execution_id: Workflow execution ID
        * table_metadata: Source TableMetadata
        * run_dt_utc: UTC timestamp object

Returns:
        * None


def sample_point(input_features: torch.Tensor, point_coordinates: torch.Tensor, add_dim=False, **kwargs) -> torch.Tensor:
    if point_coordinates.dim() == 3:
        add_dim = True
        point_coordinates = point_coordinates.unsqueeze(2)
    point_features = torch.nn.functional.grid_sample(input_features, 2.0 * point_coordinates - 1.0, **kwargs)
    if add_dim:
        point_features = point_features.squeeze(3)
    return point_features

A wrapper around `torch.nn.functional.grid_sample` to support 3D point_coordinates tensors.

Args:
        input_features (`torch.Tensor` of shape (batch_size, channels, height, width)):
            A tensor that contains features map on a height * width grid
        point_coordinates (`torch.Tensor` of shape (batch_size, num_points, 2) or (batch_size, grid_height, grid_width,:
        2)):
            A tensor that contains [0, 1] * [0, 1] normalized point coordinates
        add_dim (`bool`):
            boolean value to keep track of added dimension

Returns:
        point_features (`torch.Tensor` of shape (batch_size, channels, num_points) or (batch_size, channels,
        height_grid, width_grid):
            A tensor that contains features for points in `point_coordinates`.


def add_to_tensor(self, mat, name='add_to_tensor'):
    with self._name_scope(name):
        multiplier_vector = array_ops.expand_dims(self.multiplier, -1)
        mat = tensor_conversion.convert_to_tensor_v2_with_dispatch(mat, name='mat')
        mat_diag = array_ops.matrix_diag_part(mat)
        new_diag = multiplier_vector + mat_diag
        return array_ops.matrix_set_diag(mat, new_diag)

Add matrix represented by this operator to `mat`.  Equiv to `I + mat`.

Args:
        mat:  `Tensor` with same `dtype` and shape broadcastable to `self`.
        name:  A name to give this `Op`.

Returns:
        A `Tensor` with broadcast shape and same `dtype` as `self`.


def AddTokenToState(self, newline, dry_run, must_split=False):
    self._PushParameterListState(newline)
    penalty = 0
    if newline:
        penalty = self._AddTokenOnNewline(dry_run, must_split)
    else:
        self._AddTokenOnCurrentLine(dry_run)
    penalty += self._CalculateComprehensionState(newline)
    penalty += self._CalculateParameterListState(newline)
    return self.MoveStateToNextToken() + penalty

Add a token to the format decision state.

    Allow the heuristic to try out adding the token with and without a newline.
    Later on, the algorithm will determine which one has the lowest penalty.

Args:
        newline: (bool) Add the token on a new line if True.
        dry_run: (bool) Don't commit whitespace changes to the FormatToken if
        True.
        must_split: (bool) A newline was required before this token.

Returns:
        The penalty of splitting after the current token.


def singleprint(self) -> fingerprinting_pywrap.Singleprint:
    try:
        return fingerprinting_pywrap.Singleprint(self.graph_def_program_hash, self.signature_def_hash, self.saved_object_graph_hash, self.checkpoint_hash)
    except (TypeError, fingerprinting_pywrap.FingerprintException) as e:
        raise ValueError(f'Encounted invalid fingerprint values when constructing singleprint.graph_def_program_hash: {self.graph_def_program_hash}signature_def_hash: {self.signature_def_hash}saved_object_graph_hash: {self.saved_object_graph_hash}checkpoint_hash: {self.checkpoint_hash}{e}') from None

Canonical fingerprinting ID for a SavedModel.

    Uniquely identifies a SavedModel based on the regularized fingerprint
    attributes. (saved_model_checksum is sensitive to immaterial changes and
    thus non-deterministic.)

Returns:
        The string concatenation of `graph_def_program_hash`,
        `signature_def_hash`, `saved_object_graph_hash`, and `checkpoint_hash`
        fingerprint attributes (separated by '/').

Raises:
        ValueError: If the fingerprint fields cannot be used to construct the
        singleprint.


def get_header(graphs, proto_fileformat='rawproto', default_ops='NoOp:NoOp,_Recv:RecvOp,_Send:SendOp'):
    ops_and_kernels = get_ops_and_kernels(proto_fileformat, graphs, default_ops)
    if not ops_and_kernels:
        print('Error reading graph!')
        return 1
    return get_header_from_ops_and_kernels(ops_and_kernels, default_ops == 'all')

Computes a header for use with tensorflow SELECTIVE_REGISTRATION.

Args:
        graphs: a list of paths to GraphDef files to include.
        proto_fileformat: optional format of proto file, either 'textproto',
        'rawproto' (default) or ops_list. The ops_list is the file contain the
        list of ops in JSON format, Ex: "[["Transpose", "TransposeCpuOp"]]".
        default_ops: optional comma-separated string of operator:kernel pairs to
        always include implementation for. Pass 'all' to have all operators and
        kernels included. Default: 'NoOp:NoOp,_Recv:RecvOp,_Send:SendOp'.

Returns:
        the string of the header that should be written as ops_to_register.h.


def filter_hunks_by_path(hunks: Iterable[diff_parser.Hunk], *, path_regexes: list[str], path_regex_exclusions: list[str]) -> list[diff_parser.Hunk]:
    if not path_regexes:
        path_regexes = ['.*']
    path_regexes = [re.compile(regex) for regex in path_regexes]

    def should_include(path: str) -> bool:
        return any((regex.search(path) for regex in path_regexes))
    path_regex_exclusions = [re.compile(regex) for regex in path_regex_exclusions]

    def should_exclude(path: str) -> bool:
        return any((regex.search(path) for regex in path_regex_exclusions))
    return [hunk for hunk in hunks if should_include(hunk.file) and (not should_exclude(hunk.file))]

Filters files according to path_regexes.

    If a file matches both a path_regex and a path_regex_exclusion, then
    it will be filtered out.

Args:
        hunks: A sequence of Hunk objects representing the hunks of the diff in the
        change.
        path_regexes: A list of regexes. Paths matching these will pass through the
        filter. By default, every path is matched.
        path_regex_exclusions: A list of regexes. Paths that match both a path_regex
        and a path_regex_exclusion won't pass through the filter.

Returns:
        A list of FileDiffs whose paths match a path_regex and don't match
        any path_regex_exclusions.


def _Consumers(t, func_graphs):
    consumers = t.consumers()
    for func in func_graphs:
        for input_t, placeholder in _Captures(func):
            if input_t is t:
                consumers.extend(_Consumers(placeholder, func_graphs))
    return consumers

Returns the consumers of t, crossing closure boundaries where necessary.

Args:
        t: Tensor
        func_graphs: a list of FuncGraphs that may have captured t.

Returns:
        A list of tensors. The tensors will be from the current graph and/or
        func_graphs.


def _bigbird_block_rand_mask(self, from_seq_length, to_seq_length, from_block_size, to_block_size, num_rand_blocks, last_idx=-1):
    if from_seq_length // from_block_size != to_seq_length // to_block_size:
        raise ValueError('Error the number of blocks needs to be same!')
    rand_attn = np.zeros((from_seq_length // from_block_size - 2, num_rand_blocks), dtype=np.int32)
    if not self.training:
        return rand_attn
    middle_seq = np.arange(1, to_seq_length // to_block_size - 1, dtype=np.int32)
    last = to_seq_length // to_block_size - 1
    if last_idx > 2 * to_block_size:
        last = last_idx // to_block_size - 1
    r = num_rand_blocks
    for i in range(1, from_seq_length // from_block_size - 1):
        start = i - 2
        end = i
        if i == 1:
            rand_attn[i - 1, :] = np.random.permutation(middle_seq[2:last])[:r]
        elif i == 2:
            rand_attn[i - 1, :] = np.random.permutation(middle_seq[3:last])[:r]
        elif i == from_seq_length // from_block_size - 3:
            rand_attn[i - 1, :] = np.random.permutation(middle_seq[:last])[:r]
        elif i == from_seq_length // from_block_size - 2:
            rand_attn[i - 1, :] = np.random.permutation(middle_seq[:last])[:r]
        elif start > last:
            start = last
            rand_attn[i - 1, :] = np.random.permutation(middle_seq[:start])[:r]
        elif end + 1 == last:
            rand_attn[i - 1, :] = np.random.permutation(middle_seq[:start])[:r]
        else:
            rand_attn[i - 1, :] = np.random.permutation(np.concatenate((middle_seq[:start], middle_seq[end + 1:last])))[:r]
    return rand_attn

Create adjacency list of random attention.

Args:
        from_seq_length: int. length of from sequence.
        to_seq_length: int. length of to sequence.
        from_block_size: int. size of block in from sequence.
        to_block_size: int. size of block in to sequence.
        num_rand_blocks: int. Number of random chunks per row.
        last_idx: if -1 then num_rand_blocks blocks chosen anywhere in to sequence,
        if positive then num_rand_blocks blocks chosen only up to last_idx.

Returns:
        adjacency list of size from_seq_length//from_block_size-2 by num_rand_blocks


def _get_analytic_path(path: str, elem_id: str) -> str:
    extension_slice_names = _EXTENSION_SLICE_NAMES_RE.findall(elem_id)
    if extension_slice_names:
        extension_path_elements = _EXTENSION_PATH_ELEMENTS_RE.finditer(path)
        index_adjustment = 0
        for slice_name, path_element in zip(extension_slice_names, extension_path_elements):
            start = path_element.start(1) + index_adjustment
            end = path_element.end(1) + index_adjustment
            path = path[:start] + slice_name + path[end:]
            index_adjustment += len(slice_name) - 9
    path = path.replace('[x]', '')
    return path

Builds a usable FHIRPath given an element path and id.

    Removes choice type indicators '[x]'

    Some element definitions reference attributes from extension slices. These
    elements will have ids like 'Foo.extension:someExtension' but
    paths like 'Foo.extension' We want to treat these paths as
    'Foo.someExtension' so we replace the 'extension' part of the
    path with the slice name.

Args:
        path: The path attribute from the element definition to clean.
        elem_id: The id attribute from the element definition to clean.

Returns:
        The cleaned path name.


def Encode(string, encoding=None):
    del encoding
    return string

Encode the text string to a byte string.

Args:
        string: str, The text string to encode.
        encoding: The suggested encoding if known.

Returns:
        str, The binary string.


def get_image_features(self, pixel_values: torch.FloatTensor, pixel_attention_mask: torch.LongTensor=None):
    batch_size, num_images, num_channels, height, width = pixel_values.shape
    pixel_values = pixel_values.view(batch_size * num_images, *pixel_values.shape[2:])
    nb_values_per_image = pixel_values.shape[1:].numel()
    real_images_inds = (pixel_values == 0.0).sum(dim=(-1, -2, -3)) != nb_values_per_image
    if not any(real_images_inds):
        real_images_inds[0] = True
    pixel_values = pixel_values[real_images_inds].contiguous()
    if pixel_attention_mask is None:
        pixel_attention_mask = torch.ones(size=[pixel_values.shape[i] for i in (0, 2, 3)], dtype=torch.bool, device=pixel_values.device)
    else:
        pixel_attention_mask = pixel_attention_mask.view(batch_size * num_images, *pixel_attention_mask.shape[2:])
        pixel_attention_mask = pixel_attention_mask[real_images_inds].contiguous()
    patch_size = self.config.vision_config.patch_size
    patches_subgrid = pixel_attention_mask.unfold(dimension=1, size=patch_size, step=patch_size)
    patches_subgrid = patches_subgrid.unfold(dimension=2, size=patch_size, step=patch_size)
    patch_attention_mask = (patches_subgrid.sum(dim=(-1, -2)) > 0).bool()
    image_hidden_states = self.vision_model(pixel_values=pixel_values, patch_attention_mask=patch_attention_mask)
    image_hidden_states = image_hidden_states.last_hidden_state
    image_hidden_states = self.connector(image_hidden_states)
    return image_hidden_states

Encodes images into continuous embeddings that can be forwarded to the language model.

Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
            The tensors corresponding to the input images.
        pixel_attention_mask (`torch.LongTensor`, *optional*):
            The attention mask indicating padded regions in the image.


def register(self, obj, value):
    if obj in self._registry:
        raise KeyError(f'{type(obj)} has already been registered.')
    self._registry[obj] = value

Registers a Python object within the registry.

Args:
        obj: The object to add to the registry.
        value: The stored value for the 'obj' type.

Raises:
        KeyError: If the same obj is used twice.


def _from_config(cls, config, **kwargs):
    torch_dtype = kwargs.pop('torch_dtype', config.torch_dtype)
    if isinstance(torch_dtype, str):
        torch_dtype = getattr(torch, torch_dtype)
    dtype_orig = None
    if torch_dtype is not None:
        dtype_orig = cls._set_default_torch_dtype(torch_dtype)
    config = copy.deepcopy(config)
    if config._attn_implementation_internal is not None:
        attn_implementation = config._attn_implementation_internal
    else:
        attn_implementation = None
    config._attn_implementation = kwargs.pop('attn_implementation', attn_implementation)
    if not getattr(config, '_attn_implementation_autoset', False):
        config = cls._autoset_attn_implementation(config, check_device_map=False, torch_dtype=torch_dtype)
    if is_deepspeed_zero3_enabled() and (not _is_quantized) and (not _is_ds_init_called):
        logger.info('Detected DeepSpeed ZeRO-3: activating zero.init() for this model')
        import deepspeed
        init_contexts = [deepspeed.zero.Init(config_dict_or_path=deepspeed_config()), set_zero3_state()]
        with ContextManagers(init_contexts):
            model = cls(config, **kwargs)
    else:
        model = cls(config, **kwargs)
    if dtype_orig is not None:
        torch.set_default_dtype(dtype_orig)
    return model

All context managers that the model should be initialized under go here.

Args:
        torch_dtype (`torch.dtype`, *optional*):
            Override the default `torch.dtype` and load the model under this dtype.


def most_recent_n(self, n):
    return self._commands[-n:]

Look up the n most recent commands.

Args:
        n: Number of most recent commands to look up.

Returns:
        A list of n most recent commands, or all available most recent commands,
        if n exceeds size of the command history, in chronological order.


def __init__(self, dtype, size=None, dynamic_size=None, clear_after_read=None, tensor_array_name=None, handle=None, flow=None, infer_shape=True, element_shape=None, colocate_with_first_write_call=True, name=None):
    assert handle is None
    del handle
    del clear_after_read
    del tensor_array_name
    del colocate_with_first_write_call
    self._dynamic_size = dynamic_size
    self._size = size
    if flow is not None and (not isinstance(flow, tensor_lib.Tensor) or flow.dtype != dtypes.variant):
        raise TypeError(f'Expected `flow` to be a variant tensor, but received `{flow.dtype}` instead.')
    if flow is None and size is None:
        raise ValueError('Argument `size` must be provided if argument `flow` is not provided.')
    if flow is not None and size is not None:
        raise ValueError('Cannot provide both `flow` and `size` arguments at the same time.')
    if flow is not None and element_shape is not None:
        raise ValueError('Cannot provide both `flow` and `element_shape` argumentsat the same time.')
    self._dtype = dtypes.as_dtype(dtype).base_dtype
    self._element_shape = [tensor_shape.as_shape(element_shape)]
    self._infer_shape = infer_shape
    with ops.name_scope(name, 'TensorArrayV2', [size, flow]) as scope:
        if flow is None:
            self._flow = list_ops.tensor_list_reserve(element_shape=element_shape, num_elements=size, element_dtype=dtype, name=scope)
        else:
            self._flow = flow
    self._colocate_with_first_write_call = None
    self._colocate_with = None

Constructs a graph mode TensorArray.

Args:
        dtype: (required) data type of the TensorArray.
        size: (optional) int32 scalar `Tensor`: the size of the TensorArray.
        Required if flow is not provided.
        dynamic_size: (optional) Python bool: If true, writes to the TensorArray
        can grow the TensorArray past its initial size.  Default: False.
        clear_after_read: (optional) unused. Not supported in TensorLists.
        tensor_array_name: (optional) unused.
        handle: (optional) Must always be None.
        flow: (optional) A variant `Tensor` scalar for a TensorList.
        infer_shape: (optional, default: True) If True, shape inference is
        enabled.  In this case, all elements must have the same shape.
        element_shape: (optional, default: None) A `TensorShape` object specifying
        the shape constraints of each of the elements of the TensorArray. Need
        not be fully defined.
        colocate_with_first_write_call: (optional). unused.
        name: (optional) A name for the operation.

Raises:
        ValueError: if both handle and tensor_array_name are provided.
        TypeError: if handle is provided but is not a Tensor.


def get_num_bytes(self, batch: Sequence[tf.Tensor]) -> int:
    return sum((sys.getsizeof(element) for element in batch))

Returns:
    The number of bytes of data for a batch of Tensors.


def __init__(self, name, qual_type, external):
    Qualifier.__init__(self, name, qual_type)
    self.external_name = external
    module, event = external.rsplit('.', 1)
    self.external = importlib.import_module(module).__getattribute__(event)()

Wraps an external qualifier defined in a python file.

Example:
    qualifier int RandomInt() = external "foo.bar.RandomInt";

    There should be a python module "foo.bar" with a class RandomInt that inherits
    stl.lib.Qualifier.

Attributes:
        external_name: Name of class that extends stl.lib.Qualifier.
        external: An instance of the stl.lib.Qualifier class.


def __init__(self, window_ms: int=1, bucket_ms: int=1, overload_ratio: float=2, delay_secs: int=5):
    self.throttler = AdaptiveThrottler(window_ms=window_ms, bucket_ms=bucket_ms, overload_ratio=overload_ratio)
    self.delay_secs = delay_secs

Default throttler that uses
    :class:`apache_beam.io.components.adaptive_throttler.AdaptiveThrottler`

Args:
        window_ms (int): length of history to consider, in ms, to set throttling.
        bucket_ms (int): granularity of time buckets that we store data in, in ms.
        overload_ratio (float): the target ratio between requests sent and
        successful requests. This is "K" in the formula in
        https://landing.google.com/sre/book/chapters/handling-overload.html.
        delay_secs (int): minimum number of seconds to throttle a request.


def do_quantize_training_on_graphdef(input_graph, num_bits):
    graph = graph_pb2.GraphDef()
    result_graph_string = DoQuantizeTrainingOnGraphDefHelper(input_graph.SerializeToString(), num_bits)
    graph.ParseFromString(result_graph_string)
    return graph

A general quantization scheme is being developed in `tf.contrib.quantize`.

    Consider using that instead, though since it is in the tf.contrib namespace,
    it is not subject to backward compatibility guarantees.

Args:
        input_graph: A `GraphDef`.
        num_bits: The number of bits for quantize training.

Returns:
        The graph with quantize training done.


def run_metadata_graphs(name, data, step=None):
    summary_metadata = summary_pb2.SummaryMetadata()
    summary_metadata.plugin_data.plugin_name = 'graph_run_metadata_graph'
    summary_metadata.plugin_data.content = b'1'
    data = config_pb2.RunMetadata(function_graphs=data.function_graphs, partition_graphs=data.partition_graphs)
    with summary_scope(name, 'graph_run_metadata_graph_summary', [data, step]) as (tag, _):
        with ops.device('cpu:0'):
            tensor = constant_op.constant(data.SerializeToString(), dtype=dtypes.string)
        return write(tag=tag, tensor=tensor, step=step, metadata=summary_metadata)

Writes graphs from a RunMetadata summary.

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

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

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


def __init__(self, features, retry=None, timeout=120, max_batch_size=None, min_batch_size=None, client_options=None, metadata=None):
    super().__init__(features=features, retry=retry, timeout=timeout, max_batch_size=max_batch_size, min_batch_size=min_batch_size, client_options=client_options, metadata=metadata)

Args:
    features: (List[``vision.Feature``]) Required.
        The Vision API features to detect
    retry: (google.api_core.retry.Retry) Optional.
        A retry object used to retry requests.
        If None is specified (default), requests will not be retried.
    timeout: (float) Optional.
        The time in seconds to wait for the response from the Vision API.
        Default is 120.
    max_batch_size: (int) Optional.
        Maximum number of images to batch in the same request to the Vision API.
        Default is 5 (which is also the Vision API max).
        This parameter is primarily intended for testing.
    min_batch_size: (int) Optional.
        Minimum number of images to batch in the same request to the Vision API.
        Default is None. This parameter is primarily intended for testing.
    client_options:
        (Union[dict, google.api_core.client_options.ClientOptions]) Optional.
        Client options used to set user options on the client.
        API Endpoint should be set through client_options.
    metadata: (Optional[Sequence[Tuple[str, str]]]): Optional.
        Additional metadata that is provided to the method.


def update_state(self, y_true, y_pred, sample_weight=None):
    y_true = math_ops.cast(y_true, self._dtype)
    y_pred = math_ops.cast(y_pred, self._dtype)
    y_pred, y_true = losses_utils.squeeze_or_expand_dimensions(y_pred, y_true)
    error_sq = math_ops.squared_difference(y_pred, y_true)
    return super(RootMeanSquaredError, self).update_state(error_sq, sample_weight=sample_weight)

Accumulates root mean squared error statistics.

Args:
        y_true: The ground truth values.
        y_pred: The predicted values.
        sample_weight: Optional weighting of each example. Defaults to 1. Can be a
        `Tensor` whose rank is either 0, or the same rank as `y_true`, and must
        be broadcastable to `y_true`.

Returns:
        Update op.


def _merge_heads(self, x: torch.Tensor) -> torch.Tensor:
    batch_size_and_num_heads, seq_length, _ = x.shape
    batch_size = batch_size_and_num_heads // self.num_heads
    x = x.view(batch_size, self.num_heads, seq_length, self.head_dim)
    x = x.permute(0, 2, 1, 3)
    return x.reshape(batch_size, seq_length, self.num_heads * self.head_dim)

Merge heads together over the last dimension

Args:
        x (`torch.tensor`): [batch_size * num_heads, seq_length, head_dim]

Returns:
        torch.tensor: [batch_size, seq_length, num_heads * head_dim]


def __init__(self, step, metric, labels=None):
    self.step = step
    self.metric = metric
    self.labels = labels if labels else {}

Initializes ``MetricKey``.

Args:
        step: A string with the step this metric cell is part of.
        metric: A ``MetricName`` namespace+name that identifies a metric.
        labels: An arbitrary set of labels that also identifies the metric.


def _iceberg_io_read_test_preprocessor(test_spec: dict, expected: List[str], env: TestEnvironment):
    if (pipeline := test_spec.get('pipeline', None)):
        for transform in pipeline.get('transforms', []):
            if transform.get('type', '') == 'ReadFromIceberg':
                config = transform['config']
                db_name, table_name, field_value_dynamic_destinations = config['table'].split('.')
                transform['type'] = 'Create'
                transform['config'] = {k: v for k, v in config.items() if k.startswith('__')}
                transform['config']['elements'] = INPUT_TABLES[str(db_name), str(table_name), str(field_value_dynamic_destinations)]
    return test_spec

Preprocessor for tests that involve reading from Iceberg.

    This preprocessor replaces any ReadFromIceberg transform with a Create
    transform that reads from a predefined in-memory dictionary. This allows
    the test to verify the pipeline's correctness without relying on Iceberg
    tables stored externally.

Args:
        test_spec: The dictionary representation of the YAML pipeline specification.
        expected: A list of strings representing the expected output of the
        pipeline.
        env: The TestEnvironment object providing utilities for creating temporary
        files.

Returns:
        The modified test_spec dictionary with ReadFromIceberg transforms replaced.


def get_unconditional_inputs(self, num_samples=1, return_tensors='pt'):
    inputs = self.tokenizer([''] * num_samples, return_tensors=return_tensors, return_attention_mask=True)
    inputs['attention_mask'][:] = 0
    return inputs

Helper function to get null inputs for unconditional generation, enabling the model to be used without the
    feature extractor or tokenizer.

Args:
        num_samples (int, *optional*):
            Number of audio samples to unconditionally generate.

Example:
    ```python
    >>> from transformers import MusicgenMelodyForConditionalGeneration, MusicgenMelodyProcessor

    >>> model = MusicgenMelodyForConditionalGeneration.from_pretrained("facebook/musicgen-melody")

    >>> # get the unconditional (or 'null') inputs for the model
    >>> processor = MusicgenMelodyProcessor.from_pretrained("facebook/musicgen-melody")
    >>> unconditional_inputs = processor.get_unconditional_inputs(num_samples=1)

    >>> audio_samples = model.generate(**unconditional_inputs, max_new_tokens=256)
    ```


def _make_callable(self, feed_arrays, feed_symbols, symbol_vals, session):
    callable_opts = config_pb2.CallableOptions()
    for x in feed_arrays:
        callable_opts.feed.append(x.name)
    if self.feed_dict:
        for key in sorted(self.feed_dict.keys()):
            callable_opts.feed.append(key.name)
    for x, y in zip(feed_symbols, symbol_vals):
        connection = callable_opts.tensor_connection.add()
        if x.dtype != y.dtype:
            y = math_ops.cast(y, dtype=x.dtype)
        from_tensor = _as_graph_element(y)
        if from_tensor is None:
            from_tensor = y
        connection.from_tensor = from_tensor.name
        connection.to_tensor = x.name
    for x in self.outputs + self.fetches:
        callable_opts.fetch.append(x.name)
    callable_opts.target.append(self.updates_op.name)
    if self.run_options:
        callable_opts.run_options.CopyFrom(self.run_options)
    callable_fn = session._make_callable_from_options(callable_opts)
    self._callable_fn = callable_fn
    self._feed_arrays = feed_arrays
    self._feed_symbols = feed_symbols
    self._symbol_vals = symbol_vals
    self._fetches = list(self.fetches)
    self._session = session

Generates a callable that runs the graph.

Args:
        feed_arrays: List of input tensors to be fed Numpy arrays at runtime.
        feed_symbols: List of input tensors to be fed symbolic tensors at runtime.
        symbol_vals: List of symbolic tensors to be fed to `feed_symbols`.
        session: Session to use to generate the callable.

Returns:
        Function that runs the graph according to the above options.


def center_crop(self, image: np.ndarray, crop_size: Dict[str, int], size: Optional[int]=None, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None, **kwargs) -> np.ndarray:
    size = self.size if size is None else size
    size = get_size_dict(size)
    crop_size = get_size_dict(crop_size, param_name='crop_size')
    height, width = get_image_size(image, channel_dim=input_data_format)
    min_dim = min(height, width)
    cropped_height = size['height'] / crop_size['height'] * min_dim
    cropped_width = size['width'] / crop_size['width'] * min_dim
    return center_crop(image, size=(cropped_height, cropped_width), data_format=data_format, input_data_format=input_data_format, **kwargs)

Center crop an image to `(size["height"] / crop_size["height"] * min_dim, size["width"] / crop_size["width"] *
    min_dim)`. Where `min_dim = min(size["height"], size["width"])`.

    If the input size is smaller than `crop_size` along any edge, the image will be padded with zeros and then
    center cropped.

Args:
        image (`np.ndarray`):
            Image to center crop.
        crop_size (`Dict[str, int]`):
            Desired output size after applying the center crop.
        size (`Dict[str, int]`, *optional*):
            Size of the image after resizing. If not provided, the self.size attribute will be used.
        data_format (`str` or `ChannelDimension`, *optional*):
            The channel dimension format of the image. If not provided, it will be the same as the input image.
        input_data_format (`str` or `ChannelDimension`, *optional*):
            The channel dimension format of the input image. If not provided, it will be inferred.


def _static_value_provider_of(value_type):

    def _f(value):
        _f.__name__ = value_type.__name__
        return StaticValueProvider(value_type, value)
    return _f

Helper function to plug a ValueProvider into argparse.

Args:
        value_type: the type of the value. Since the type param of argparse's
                add_argument will always be ValueProvider, we need to
                preserve the type of the actual value.

Returns:
        A partially constructed StaticValueProvider in the form of a function.


def __init__(self, num_groups=2):
    if num_groups < 1:
        raise ValueError(f'Argument `num_groups` must be a positive integer. Received: num_groups={num_groups}')
    self._ready = threading.Condition(threading.Lock())
    self._num_groups = num_groups
    self._group_member_counts = [0] * self._num_groups

Initialize a group lock.

Args:
        num_groups: The number of groups that will be accessing the resource under
        consideration. Should be a positive number.

Returns:
        A group lock that can then be used to synchronize code.

Raises:
        ValueError: If num_groups is less than 1.


def _preprocess_async_for_and_yield(idx: int, bytecode: Sequence[opcodes.Opcode], prev_block: Block, processed_blocks: set[Block]) -> tuple[list[Block], int]:
    assert isinstance(bytecode[idx], opcodes.SEND)
    i = next((i for i in range(idx + 1, len(bytecode)) if isinstance(bytecode[i], opcodes.JUMP_BACKWARD_NO_INTERRUPT)))
    end_block_idx = i + 1
    if isinstance(bytecode[end_block_idx], opcodes.CLEANUP_THROW):
        end_block_idx += 1
    send_block = Block(bytecode[idx:idx + 1])
    yield_value_block = Block(bytecode[idx + 1:end_block_idx])
    prev_block.connect_outgoing(send_block)
    send_block.connect_outgoing(yield_value_block)
    processed_blocks.update(send_block, yield_value_block)
    return ([send_block, yield_value_block], end_block_idx)

Process bytecode instructions for yield and async for in a way that pytype can iterate correctly.

    'Async for' and yield statements, contains instructions that starts with SEND
    and ends with END_SEND.

    The reason why we need to pre process async for is because the control flow of
    async for is drastically different from regular control flows also due to the
    fact that the termination of the loop happens by STOP_ASYNC_ITERATION
    exception, not a regular control flow. So we need to split (or merge) the
    basic blocks in a way that pytype executes in the order that what'd happen in
    the runtime, so that it doesn't fail with wrong order of execution, which can
    result in a stack underrun.

Args:
        idx: The index of the SEND instruction.
        bytecode: A list of instances of opcodes.Opcode
        prev_block: The previous block that we want to connect the new blocks to.
        processed_blocks: Blocks that has been processed so that it doesn't get
        processed again by compute_order.

Returns:
        A tuple of (list[Block], int), where the Block is the block containing the
        iteration part of the async for construct, and the int is the index of the
        END_SEND instruction.


def __init__(self, schema: dict[str, Any], error_handling: Optional[Mapping[str, Any]]=None):
    self._schema = schema
    self._exception_handling_args = exception_handling_args(error_handling)

Validates each element of a PCollection against a json schema.

Args:
        schema: A json schema against which to validate each element.
        error_handling: Whether and how to handle errors during iteration.
            If this is not set, invalid elements will fail the pipeline, otherwise
            invalid elements will be passed to the specified error output along
            with information about how the schema was invalidated.


def _strip_padding(tensors, pad_len):
    if not tensors:
        raise ValueError('tensors cannot be empty')
    shape = tensors[0].shape
    if len(shape) > 1:
        raise ValueError('tensors must be 1D')
    prefix_len = int(shape[0] - pad_len)
    if prefix_len < 0:
        raise ValueError('pad_len longer than tensor')
    stripped = []
    for t in tensors:
        with ops.colocate_with(t):
            stripped.append(array_ops.slice(t, [0], [prefix_len]))
    return stripped

Strip the suffix padding added by _padded_split.

Args:
        tensors: list of `tf.Tensor` of identical length 1D tensors.
        pad_len: number of elements to be stripped from the end of each tensor.

Returns:
        list of `tf.Tensor` which are the stripped inputs.

Raises:
        ValueError: tensors must be a non-empty list of 1D tensors, and
        each must be longer than pad_len.


def __init__(self, constructor: Callable[[], T], tag: Any, *, path: str=tempfile.gettempdir(), always_proxy: Optional[bool]=None):
    self._constructor = constructor
    self._tag = tag
    self._path = path
    self._always_proxy = False if always_proxy is None else always_proxy
    self._proxy = None
    self._manager = None
    self._rpc_address = None
    self._cross_process_lock = fasteners.InterProcessLock(os.path.join(self._path, self._tag) + '.lock')

MultiProcessShared is used to share a single object across processes.

    For example, one could have the class::

        class MyExpensiveObject(object):
        def __init__(self, args):
        [expensive initialization and memory allocation]

        def method(self, arg):
        ...

    One could share a single instance of this class by wrapping it as::

        shared_ptr = MultiProcessShared(lambda: MyExpensiveObject(...))
        my_expensive_object = shared_ptr.acquire()

    which could then be invoked as::

        my_expensive_object.method(arg)

    This can then be released with::

        shared_ptr.release(my_expensive_object)

    but care should be taken to avoid releasing the object too soon or
    expensive re-initialization may be required, defeating the point of
    using a shared object.

Args:
        constructor: function that initialises / constructs the object if not
        present in the cache. This function should take no arguments. It should
        return an initialised object, or raise an exception if the object could
        not be initialised / constructed.
        tag: an indentifier to store with the cached object. If multiple
        MultiProcessShared instances are created with the same tag, they will all
        share the same proxied object.
        path: a temporary path in which to create the inter-process lock
        always_proxy: whether to direct all calls through the proxy, rather than
        call the object directly for the process that created it


def argmin(x, axis=-1):
    return math_ops.argmin(x, axis)

Returns the index of the minimum value along an axis.

Args:
        x: Tensor or variable.
        axis: axis along which to perform the reduction.

Returns:
        A tensor.


def __init__(self, config, in_features, condition_dim, n_classes=256, bottleneck_factor=2):
    super().__init__()
    bottleneck = (in_features + condition_dim) // bottleneck_factor
    self.mlp = nn.Sequential(nn.Conv2d(in_features + condition_dim, bottleneck, kernel_size=1, stride=1, padding=0), nn.GELU(), nn.Conv2d(bottleneck, 2 + 2, kernel_size=1, stride=1, padding=0), nn.Softplus())
    self.p_eps = 0.0001
    self.max_temp = config.max_temp
    self.min_temp = config.min_temp
    self.log_binomial_transform = LogBinomialSoftmax(n_classes, act=torch.softmax)

Per-pixel MLP followed by a Conditional Log Binomial softmax.

Args:
        in_features (`int`):
            Number of input channels in the main feature.
        condition_dim (`int`):
            Number of input channels in the condition feature.
        n_classes (`int`, *optional*, defaults to 256):
            Number of classes.
        bottleneck_factor (`int`, *optional*, defaults to 2):
            Hidden dim factor.


def relayout(tensor: tensor_lib.Tensor, layout: layout_lib.Layout, name: Optional[str]=None) -> tensor_lib.Tensor:
    layout_str = layout.to_string()
    with default_mesh(layout.mesh):
        return gen_dtensor_ops.relayout(tensor, layout_str, name=name)

Changes the layout of `tensor`.

    Changes the layout of `tensor` to `layout`. This is used to fine-tune the
    behavior of ops following/connected to `tensor`, such as choosing one SPMD
    expansion pattern over another. This works by forward propagating `layout`
    to connected TensorFlow computation graphs during layout propagation.

    Currently, only converting layouts from replicated to sharded or sharded to
    replicated per mesh dimension is supported. That is, "x, y" -> "unsharded, y"
    is supported, while "x, y" -> "z, y" is not supported.

    We also support a special "match" sharding spec, which instructs the relayout
    to act as an identity operation with respect to any sharding on these
    mesh dimensions.

    Relayout is internally lowered to a set of Split and/or AllToAll ops. When
    tensor layouts are converted from replicated to sharded, the cost is
    comparatively low because we only insert Split ops and no cross-device
    communication is needed. However, when tensor layouts are converted from
    sharded to replicated, cross-device communication may occur, causing potential
    performance impact.

Args:
        tensor: A DTensor to specify a new layout for.
        layout: A Layout object specifying a new sharding spec.
        name: name of the Op.

Returns:
        A DTensor output from the Relayout op.


def cancel(self):
    raise NotImplementedError()

Cancels the pipeline execution.

Raises:
        IOError: If there is a persistent problem getting job
        information.
        NotImplementedError: If the runner does not support this
        operation.

Returns:
        The final state of the pipeline.


def reduce_mean(values, index, name='segmented_reduce_mean'):
    return _segment_reduce(values, index, tf.math.unsorted_segment_mean, name)

Averages a tensor over its segments. Outputs 0 for empty segments. This operations computes the mean over segments,
    with support for:

        - Batching using the first dimensions [B1, B2, ..., Bn]. Each element in a batch can have different indices.
        - Vectorization using the last dimension [V1, V2, ...]. If they are present the output will be a mean of vectors
        rather than scalars.
    Only the middle dimensions [I1, ..., Ik] are reduced by the operation.

Args:
        values: [B1, B2, ..., Bn, I1, .., Ik, V1, V2, ..] tensor of values to be
        averaged.
        index: IndexMap [B1, B2, ..., Bn, I1, .., Ik] index defining the segments.
        name: Name for the TensorFlow ops.

Returns:
        A pair (output_values, output_index) where `output_values` is a tensor of shape [B1, B2, ..., Bn, num_segments,
        V1, V2, ..] and `index` is an IndexMap with shape [B1, B2, ..., Bn, num_segments].


def num_gpus():
    return context().num_gpus()

Get the number of available GPU devices.

Returns:
        The number of available GPU devices.


def sanity_check_type(self, other):
    if type(self) is not type(other):
        raise ValueError('No TypeSpec is compatible with both %s and %s' % (self, other))
    if self._input_workers.serialize() != other._input_workers.serialize():
        raise ValueError('_input_workers is not compatible with both %s and %s' % (self, other))
    if self._strategy is not other._strategy:
        raise ValueError('tf.distribute strategy is not compatible with both %s and %s' % (self, other))

Returns the most specific TypeSpec compatible with `self` and `other`.

Args:
        other: A `TypeSpec`.

Raises:
        ValueError: If there is no TypeSpec that is compatible with both `self`
        and `other`.


def test_sample_sabr(self, supply_grad_vol_fn):
    dtype = np.float64
    drift_fn = lambda _, x: tf.zeros_like(x)
    beta = tf.constant(0.5, dtype=dtype)
    volvol = tf.constant(1.0, dtype=dtype)
    rho = tf.constant(0.2, dtype=dtype)

    def vol_fn(t, x):
        """The volatility function for the SABR model."""
        del t
        f = x[..., 0]
        v = x[..., 1]
        fb = f ** beta
        m11 = v * fb * tf.math.sqrt(1 - tf.square(rho))
        m12 = v * fb * rho
        m21 = tf.zeros_like(m11)
        m22 = volvol * v
        mc1 = tf.concat([tf.expand_dims(m11, -1), tf.expand_dims(m21, -1)], -1)
        mc2 = tf.concat([tf.expand_dims(m12, -1), tf.expand_dims(m22, -1)], -1)
        should_be_zero = tf.expand_dims(tf.expand_dims((beta != 0) & (f <= 0.0), -1), -1)
        vol_matrix = tf.concat([tf.expand_dims(mc1, -1), tf.expand_dims(mc2, -1)], -1)
        return tf.where(should_be_zero, tf.zeros_like(vol_matrix), vol_matrix)
    if supply_grad_vol_fn:

        def _grad_volatility_fn(current_time, current_state, input_gradients):
            return gradient.fwd_gradient(functools.partial(vol_fn, current_time), current_state, input_gradients=input_gradients, unconnected_gradients=tf.UnconnectedGradients.ZERO)
        grad_volatility_fn = _grad_volatility_fn
    else:
        grad_volatility_fn = None
    times = np.array([0.0, 0.1, 0.21, 0.32, 0.43, 0.55])
    x0 = np.array([0.1, 0.2])
    paths = self.evaluate(milstein_sampling.sample(dim=2, drift_fn=drift_fn, volatility_fn=vol_fn, times=times, num_samples=1000, initial_state=x0, grad_volatility_fn=grad_volatility_fn, random_type=tff.math.random.RandomType.STATELESS_ANTITHETIC, time_step=0.01, seed=[1, 42]))
    mean = np.average(paths)
    stddev = np.std(paths)
    euler_paths = self.evaluate(euler_sampling.sample(dim=2, drift_fn=drift_fn, volatility_fn=vol_fn, times=times, time_step=0.01, num_samples=10000, initial_state=x0, random_type=tff.math.random.RandomType.STATELESS_ANTITHETIC, seed=[1, 42]))
    euler_mean = np.average(euler_paths)
    euler_stddev = np.std(euler_paths)
    self.assertAllClose((mean, stddev), (euler_mean, euler_stddev), rtol=0.05, atol=0.05)

Tests path properties for SABR.

    We construct the following Ito process.

    ```
        dF_t = v_t * F_t ^ beta * dW_{F,t}
        dv_t = volvol * v_t * dW_{v,t}
        dW_{F,t} * dW_{v,t} = rho * dt
    ```

    `F_t` is the forward. `v_t` is volatility. `beta` is the CEV parameter.
    `volvol` is volatility of volatility. `W_{F,t}` and `W_{v,t}` are two
    correlated Wiener processes with instantaneous correlation `rho`.

Args:
        supply_grad_vol_fn: A bool. Whether or not to supply a grad_volatility_fn.


def annotate_source(dump, source_file_path, do_dumped_tensors=False, file_stack_top=False, min_line=None, max_line=None):
    py_graph = dump.python_graph
    if not py_graph:
        raise ValueError('Cannot perform source annotation due to a lack of set Python graph in the dump object')
    source_file_path = _norm_abs_path(source_file_path)
    line_to_op_names = {}
    for op in py_graph.get_operations():
        for file_path, line_number, _, _ in reversed(dump.node_traceback(op.name)):
            if min_line is not None and line_number < min_line or (max_line is not None and line_number >= max_line):
                continue
            if _norm_abs_path(file_path) != source_file_path:
                continue
            if do_dumped_tensors:
                watch_keys = dump.debug_watch_keys(op.name)
                items_to_append = list(set(map(_convert_watch_key_to_tensor_name, watch_keys)))
            else:
                items_to_append = [op.name]
            if line_number in line_to_op_names:
                line_to_op_names[line_number].extend(items_to_append)
            else:
                line_to_op_names[line_number] = items_to_append
            if file_stack_top:
                break
    return line_to_op_names

Annotate a Python source file with a list of ops created at each line.

    (The annotation doesn't change the source file itself.)

Args:
        dump: (`DebugDumpDir`) A `DebugDumpDir` object of which the Python graph
        has been loaded.
        source_file_path: (`str`) Path to the source file being annotated.
        do_dumped_tensors: (`str`) Whether dumped Tensors, instead of ops are to be
        used to annotate the source file.
        file_stack_top: (`bool`) Whether only the top stack trace in the
        specified source file is to be annotated.
        min_line: (`None` or `int`) The 1-based line to start annotate the source
        file from (inclusive).
        max_line: (`None` or `int`) The 1-based line number to end the annotation
        at (exclusive).

Returns:
        A `dict` mapping 1-based line number to a list of op name(s) created at
        that line, or tensor names if `do_dumped_tensors` is True.

Raises:
        ValueError: If the dump object does not have a Python graph set.


def push(self, x):
    self._queue.append(x)

Adds a new value to the data window.

Args:
        x: The value to be added to the window.


def _AssertProtoDictEquals(self, expected_dict, actual_dict, verbose=False, update_goldens=False, additional_missing_object_message='', api_version=2):
    diffs = []
    verbose_diffs = []
    expected_keys = set(expected_dict.keys())
    actual_keys = set(actual_dict.keys())
    only_in_expected = expected_keys - actual_keys
    only_in_actual = actual_keys - expected_keys
    all_keys = expected_keys | actual_keys
    updated_keys = []
    for key in all_keys:
        diff_message = ''
        verbose_diff_message = ''
        if key in only_in_expected:
            diff_message = 'Object %s expected but not found (removed). %s' % (key, additional_missing_object_message)
            verbose_diff_message = diff_message
        elif key in only_in_actual:
            diff_message = 'New object %s found (added).' % key
            verbose_diff_message = diff_message
        else:
            self.maxDiff = None
            try:
                self.assertProtoEquals(expected_dict[key], actual_dict[key])
            except AssertionError as e:
                updated_keys.append(key)
                diff_message = 'Change detected in python object: %s.' % key
                verbose_diff_message = str(e)
        if diff_message:
            diffs.append(diff_message)
            verbose_diffs.append(verbose_diff_message)
    if diffs:
        diff_count = len(diffs)
        logging.error(self._test_readme_message)
        logging.error('%d differences found between API and golden.', diff_count)
        if update_goldens:
            logging.warning(self._update_golden_warning)
            for key in only_in_expected:
                filepath = _KeyToFilePath(key, api_version)
                file_io.delete_file(filepath)
            for key in only_in_actual | set(updated_keys):
                filepath = _KeyToFilePath(key, api_version)
                file_io.write_string_to_file(filepath, text_format.MessageToString(actual_dict[key]))
        else:
            for d, verbose_d in zip(diffs, verbose_diffs):
                logging.error('    %s', d)
                logging.error('    %s', verbose_d)
            self.fail('%d differences found between API and golden.' % diff_count)
    else:
        logging.info('No differences found between API and golden.')

Diff given dicts of protobufs and report differences a readable way.

Args:
        expected_dict: a dict of TFAPIObject protos constructed from golden files.
        actual_dict: a dict of TFAPIObject protos constructed by reading from the
        TF package linked to the test.
        verbose: Whether to log the full diffs, or simply report which files were
        different.
        update_goldens: Whether to update goldens when there are diffs found.
        additional_missing_object_message: Message to print when a symbol is
        missing.
        api_version: TensorFlow API version to test.


def start_app_and_connect(self):
    try:
        self._start_app_and_connect()
    except AppStartPreCheckError:
        raise
    except Exception as e:
        self._ad.log.exception('Failed to start app and connect.')
        try:
            self.stop_app()
        except Exception:
            self._ad.log.exception('Failed to stop app after failure to start and connect.')
        raise e

Starts snippet apk on the device and connects to it.

    This wraps the main logic with safe handling

Raises:
        AppStartPreCheckError, when pre-launch checks fail.


def _run_graph(self, device, dtype, data_format, input_shape, filter_shape, strides, padding, num_iters, warmup_iters):
    graph = ops.Graph()
    with graph.as_default():
        warmup_outputs, outputs = build_graph(device, dtype, data_format, input_shape, filter_shape, strides, padding, num_iters, warmup_iters)
        config = config_pb2.ConfigProto()
        config.graph_options.optimizer_options.opt_level = -1
        rewrite_options = config.graph_options.rewrite_options
        rewrite_options.layout_optimizer = rewriter_config_pb2.RewriterConfig.ON if FLAGS.enable_layout_optimizer else rewriter_config_pb2.RewriterConfig.OFF
        rewrite_options.dependency_optimization = rewriter_config_pb2.RewriterConfig.OFF
        with session_lib.Session(graph=graph, config=config) as session:
            variables.global_variables_initializer().run()
            session.run(warmup_outputs)
            start_time = time.time()
            session.run(outputs)
            duration = (time.time() - start_time) / num_iters
            print('%s %s %s inputshape:%s filtershape:%s strides:%s padding:%s %d iters: %.8f sec' % (device, str(dtype), data_format, str(input_shape).replace(' ', ''), str(filter_shape).replace(' ', ''), str(strides).replace(' ', ''), padding, num_iters, duration))
    name_template = 'conv2d_{device}_{datatype}_{data_format}_input_shape_{inputshape}_filter_shape_{filtershape}_strides_{strides}_padding_{padding}'
    self.report_benchmark(name=name_template.format(device=device, datatype=str(dtype), data_format=str(data_format), inputshape=str(input_shape).replace(' ', ''), filtershape=str(filter_shape).replace(' ', ''), strides=str(strides).replace(' ', ''), padding=padding).replace(' ', ''), iters=num_iters, wall_time=duration)
    return duration

runs the graph and print its execution time.

Args:
        device: String, the device to run on.
        dtype: Data type for the convolution.
        data_format: A string from: "NHWC" or "NCHW". Data format for input and
                output data.
        input_shape: Shape of the input tensor.
        filter_shape: Shape of the filter tensor.
        strides: A list of ints. 1-D of length 4. The stride of sliding
            window for each dimension of input.
        padding: A string from: "SAME", "VALID". The type of padding
            algorithm to use.  num_iters: Number of iterations to run the
                benchmark.
        num_iters: number of iterations to run conv2d.
        warmup_iters: number of iterations for warmup runs.

Returns:
        The duration of the run in seconds.


def n_gpu(self):
    requires_backends(self, ['torch'])
    if not hasattr(self, '_n_gpu'):
        _ = self._setup_devices
    return self._n_gpu

The number of GPUs used by this process.

Note:
        This will only be greater than one when you have multiple GPUs available but are not using distributed
        training. For distributed training, it will always be 1.


def model_type_scope(value):
    previous_value = _thread_local_data.model_type
    try:
        _thread_local_data.model_type = value
        yield value
    finally:
        _thread_local_data.model_type = previous_value

Provides a scope within which the model type to test is equal to `value`.

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

Args:
        value: model type value

Yields:
        The provided value.


def __init__(self, drop_prob, **kwargs):
    super().__init__(**kwargs)
    self.drop_prob = drop_prob

Optimized dropout module for stabilizing the training

Args:
        drop_prob (float): the dropout probabilities


def _process_landing_page(self, item, feed_item):
    if feed_item.get(FieldMap.AD_LANDING_PAGE_ID, ''):
        landing_page = self._landing_page_dao.get(feed_item, required=True)
        item['clickThroughUrl'] = {'landingPageId': landing_page['id']}
    if feed_item.get(FieldMap.AD_URL_SUFFIX, ''):
        item['clickThroughUrlSuffixProperties'] = {'overrideInheritedSuffix': True, 'clickThroughUrlSuffix': feed_item.get(FieldMap.AD_URL_SUFFIX, '')}
    else:
        item['clickThroughUrlSuffixProperties'] = {'overrideInheritedSuffix': False}

Configures ad landing page.

Args:
        item: DCM ad object to update.
        feed_item: Feed item representing the ad from the Bulkdozer feed


def broadcast_to(x, shape):
    if any_symbolic_tensors((x,)):
        return BroadcastTo(shape=shape).symbolic_call(x)
    return backend.numpy.broadcast_to(x, shape)

Broadcast a tensor to a new shape.

Args:
        x: The tensor to broadcast.
        shape: The shape of the desired tensor. A single integer `i` is
            interpreted as `(i,)`.

Returns:
        A tensor with the desired shape.

Example:
    >>> x = keras.ops.array([1, 2, 3])
    >>> keras.ops.broadcast_to(x, (3, 3))
    array([[1, 2, 3],
        [1, 2, 3],
        [1, 2, 3]])


def to_parquet(evset: EventSet, path: str, **kwargs):
    df = to_pandas(evset)
    df.to_parquet(path, **kwargs)

Saves an [`EventSet`][temporian.EventSet] to a CSV file.

Example:
        ```python
        >>> output_path = str(tmp_dir / "output_data.parquet")
        >>> evset = tp.event_set(timestamps=[1,], features={"f1": [0.1]})
        >>> tp.to_parquet(evset, output_path)

        ```

Args:
        evset: EventSet to save.
        path: Path to the file.


def reset(self, state):
    state = _convert_to_state_tensor(state)
    state.shape.assert_is_compatible_with([_get_state_size(self.algorithm)])
    self._state_var.assign(state)

Resets the generator by a new state.

    See `__init__` for the meaning of "state".

Args:
        state: the new state.


def tanh_shrink(x):
    return ops.tanh_shrink(x)

Tanh shrink activation function.

    It is defined as:

    `f(x) = x - tanh(x)`.

Args:
        x: Input tensor.


def __sub__(self, other: Union[None, int, str, 'KeyPath']) -> 'KeyPath':
    if other is None:
        return self
    if isinstance(other, str):
        other = KeyPath.parse(other)
    elif isinstance(other, int):
        other = KeyPath(other)
    if not isinstance(other, KeyPath):
        raise TypeError(f'Cannot subtract KeyPath({self}) by {other!r}.')
    max_len = max(len(self), len(other))
    for pos in range(max_len):
        if pos >= len(self):
            raise ValueError(f'KeyPath subtraction failed: left path {self!r} is an ancestor of right path {other!r}.')
        if pos >= len(other):
            return KeyPath(self.keys[pos:])
        if self.keys[pos] != other.keys[pos]:
            raise ValueError(f'KeyPath subtraction failed: left path {self!r} and right path {other!r} are in different subtree.')
    return KeyPath()

Finds the relative path of this path to the other.

    Example::

        path1 = pg.KeyPath.parse('a.b.c.d')
        path2 = pg.KeyPath.parse('a.b')
        assert path1 - path2 == 'c.d'

Args:
        other: Object to subtract, which can be None, int (as a depth-1 KeyPath),
        string (parsed as a KeyPath) or a KeyPath object.

Returns:
        Relative path of this path to the other.

Raises:
        ValueError: This path is an ancestor node of the other path,
        or these two paths are in different branch.


def get_attr_location(self, name, location):
    line, _ = location
    src_line = self.line(line)
    attr = name.split('.')[-1]
    dot_attr = '.' + attr
    if dot_attr in src_line:
        col = src_line.index(dot_attr)
        return (Location(line, col + 1), len(attr))
    else:
        attr_loc = self._get_multiline_location(location, 5, dot_attr)
        if attr_loc:
            return (Location(attr_loc.line, attr_loc.column + 1), len(attr))
        else:
            for l in self.get_closest_line_range(line, line + 5):
                if self.line(l).endswith('.'):
                    next_line = self.next_non_comment_line(l)
                    text = self.line(next_line)
                    if text.lstrip().startswith(attr):
                        c = text.index(attr)
                        return (Location(next_line, c), len(attr))
        return (location, len(name))

Returns the location and span of the attribute in an attribute access.

Args:
        name: The attribute name.
        location: The location of the value the attribute is accessed on.


def set_learning_phase(value):
    warnings.warn('`tf.keras.backend.set_learning_phase` is deprecated and will be removed after 2020-10-11. To update it, simply pass a True/False value to the `training` argument of the `__call__` method of your layer or model.')
    deprecated_internal_set_learning_phase(value)

Sets the learning phase to a fixed value.

    The backend learning phase affects any code that calls
    `backend.learning_phase()`
    In particular, all Keras built-in layers use the learning phase as the default
    for the `training` arg to `Layer.__call__`.

    User-written layers and models can achieve the same behavior with code that
    looks like:

    ```python
        def call(self, inputs, training=None):
        if training is None:
        training = backend.learning_phase()
    ```

Args:
        value: Learning phase value, either 0 or 1 (integers).
            0 = test, 1 = train

Raises:
        ValueError: if `value` is neither `0` nor `1`.


def save_prefixed_metrics(results, output_dir, file_name: str='all_results.json', metric_key_prefix: str='eval'):
    for key in list(results.keys()):
        if not key.startswith(f'{metric_key_prefix}_'):
            results[f'{metric_key_prefix}_{key}'] = results.pop(key)
    with open(os.path.join(output_dir, file_name), 'w') as f:
        json.dump(results, f, indent=4)

Save results while prefixing metric names.

Args:
        results: (:obj:`dict`):
            A dictionary of results.
        output_dir: (:obj:`str`):
            An output directory.
        file_name: (:obj:`str`, `optional`, defaults to :obj:`all_results.json`):
            An output file name.
        metric_key_prefix: (:obj:`str`, `optional`, defaults to :obj:`eval`):
            A metric name prefix.


def __init__(self, max_length: int, eos_token_id: Union[int, List[int], torch.Tensor], device: str='cpu'):
    self.max_length = max_length
    if not isinstance(eos_token_id, torch.Tensor):
        if isinstance(eos_token_id, int):
            eos_token_id = [eos_token_id]
        eos_token_id = torch.tensor(eos_token_id, device=device)
    self.eos_token_id = eos_token_id
    if torch.is_floating_point(eos_token_id) or (eos_token_id < 0).any():
        raise ValueError(f'`eos_token_id` has to be a list of positive integers, but is {eos_token_id}')

[`LogitsProcessor`] that enforces the specified token as the last generated token when `max_length` is reached.

Args:
        max_length (`int`):
            The maximum length of the sequence to be generated.
        eos_token_id (`Union[int, List[int], torch.Tensor]`):
            The id(s) of the *end-of-sequence* token.
        device (`str`, *optional*, defaults to `"cpu"`):
            The device to allocate the tensors.

Example:
    ```python
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM

    >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2")
    >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2")

    >>> inputs = tokenizer("A sequence: 1, 2, 3", return_tensors="pt")

    >>> # By default, it continues generating according to the model's logits
    >>> outputs = model.generate(**inputs, max_new_tokens=10)
    >>> print(tokenizer.batch_decode(outputs)[0])
    A sequence: 1, 2, 3, 4, 5, 6, 7, 8

    >>> # `forced_eos_token_id` ensures the generation ends with a EOS token
    >>> outputs = model.generate(**inputs, max_new_tokens=10, forced_eos_token_id=tokenizer.eos_token_id)
    >>> print(tokenizer.batch_decode(outputs)[0])
    A sequence: 1, 2, 3, 4, 5, 6, 7,<|endoftext|>
    ```


def generator_next_fn(iterator_id_t):
    if output_types and output_shapes:
        flattened_types = [dtypes.as_dtype(dt) for dt in nest.flatten(output_types)]
        flattened_shapes = nest.flatten(output_shapes)

        def generator_py_func(iterator_id):
            """A `py_func` that will be called to invoke the iterator."""
            values = next(generator_state.get_iterator(iterator_id))
            try:
                flattened_values = nest.flatten_up_to(output_types, values)
            except (TypeError, ValueError) as e:
                raise TypeError(f'`generator` yielded an element that did not match the expected structure. The expected structure was {output_types}, but the yielded element was {values}.') from e
            ret_arrays = []
            for ret, dtype in zip(flattened_values, flattened_types):
                try:
                    ret_arrays.append(script_ops.FuncRegistry._convert(ret, dtype=dtype.as_numpy_dtype))
                except (TypeError, ValueError) as e:
                    raise TypeError(f'`generator` yielded an element that could not be converted to the expected type. The expected type was {dtype.name}, but the yielded element was {ret}.') from e
            for ret_array, expected_dtype, expected_shape in zip(ret_arrays, flattened_types, flattened_shapes):
                if ret_array.dtype != expected_dtype.as_numpy_dtype:
                    raise TypeError(f'`generator` yielded an element of type {ret_array.dtype} where an element of type {expected_dtype.as_numpy_dtype} was expected.')
                if not expected_shape.is_compatible_with(ret_array.shape):
                    raise TypeError(f'`generator` yielded an element of shape {ret_array.shape} where an element of shape {expected_shape} was expected.')
            return ret_arrays
        flat_values = script_ops.numpy_function(generator_py_func, [iterator_id_t], flattened_types)
        if not isinstance(flat_values, (list, tuple)):
            flat_values = [flat_values]
        if output_shapes is not None:
            for ret_t, shape in zip(flat_values, flattened_shapes):
                ret_t.set_shape(shape)
        return nest.pack_sequence_as(output_types, flat_values)
    else:
        flat_output_types = structure.get_flat_tensor_types(output_signature)

        def generator_py_func(iterator_id):
            """A `py_func` that will be called to invoke the iterator."""
            values = next(generator_state.get_iterator(iterator_id.numpy()))
            try:
                values = structure.normalize_element(values, output_signature)
            except (TypeError, ValueError) as e:
                raise TypeError(f'`generator` yielded an element that did not match the expected structure. The expected structure was {output_signature}, but the yielded element was {values}.') from e
            values_spec = structure.type_spec_from_value(values)
            if not structure.are_compatible(values_spec, output_signature):
                raise TypeError(f'`generator` yielded an element of {values_spec} where an element of {output_signature} was expected.')
            return structure.to_tensor_list(output_signature, values)
        return script_ops.eager_py_func(generator_py_func, inp=[iterator_id_t], Tout=flat_output_types)

Generates the next element from iterator with ID `iterator_id_t`.

    We map this function across an infinite repetition of the
    `iterator_id_t`, and raise `StopIteration` to terminate the iteration.

Args:
        iterator_id_t: A `tf.int64` tensor whose value uniquely identifies the
        iterator in `generator_state` from which to generate an element.

Returns:
        The next element to generate from the iterator.


def __init__(self, value: T, proxy: Optional[T]=None):
    if proxy is None:
        proxy = value
    super().__init__('constant', proxy)
    self._value = value

Initialize a constant expression.

Args:
        value: The constant value to be produced by this expression.
        proxy: (Optional) a proxy object with same type as `value` to use for
        rapid type checking at pipeline construction time. If not provided,
        `value` will be used directly.


def assert_no_new_tensors(f: _F) -> _F:

    def decorator(self: 'TensorFlowTestCase', **kwargs):
        """Finds existing Tensors, runs the test, checks for new Tensors."""

        def _is_tensorflow_object(obj) -> bool:
            try:
                return isinstance(obj, (tensor_lib.Tensor, variables.Variable, tensor_shape.Dimension, tensor_shape.TensorShape))
            except (ReferenceError, AttributeError):
                return False
        tensors_before = set((id(obj) for obj in gc.get_objects() if _is_tensorflow_object(obj)))
        outside_executed_eagerly = cast(bool, context.executing_eagerly())
        outside_graph_key = ops.get_default_graph()._graph_key
        with ops.Graph().as_default():
            ops.get_default_graph()._graph_key = outside_graph_key
            if outside_executed_eagerly:
                with context.eager_mode():
                    result = f(self, **kwargs)
            else:
                result = f(self, **kwargs)
        context.context()._clear_caches()
        gc.collect()
        tensors_after = [obj for obj in gc.get_objects() if _is_tensorflow_object(obj) and id(obj) not in tensors_before]
        if tensors_after:
            raise AssertionError('%d Tensors not deallocated after test: %s' % (len(tensors_after), str(tensors_after)))
        return result
    return tf_decorator.make_decorator(f, decorator)

Decorator for asserting that no new Tensors persist after a test.

    Mainly useful for checking that code using the Python C API has correctly
    manipulated reference counts.

    Clears the caches that it knows about, runs the garbage collector, then checks
    that there are no Tensor or Tensor-like objects still around. This includes
    Tensors to which something still has a reference (e.g. from missing
    Py_DECREFs) and uncollectable cycles (i.e. Python reference cycles where one
    of the objects has __del__ defined).

Args:
        f: The test case to run.

Returns:
        The decorated test case.


def testSaveAndLoadSingleVariable(self, shard_config):
    strategy = self._create_strategy(shard_config[0])
    with strategy.scope():
        var = variables_lib.Variable([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
    model_dir = self.get_temp_dir()
    save.save(var, model_dir)
    strategy2 = self._create_strategy(shard_config[1])
    with strategy2.scope():
        loaded = load.load(model_dir)
    if shard_config[1] > 1:
        loaded = array_ops.concat(loaded.variables, axis=0)
    self.assertLen(loaded.numpy(), 6)
    if shard_config[0] > 1:
        var = array_ops.concat(var.variables, axis=0)
    self.assertAllClose(var.numpy(), loaded.numpy())

Test saving and loading ShardedVariable with different numbers of shards.

    Loading tf.Variables into multiple Shards is not yet supported

Args:
        shard_config: The number of shards to use before and after loading. For
        example, [2, 1] means to create and save the variable with 2 shards and
        load it into 1 shard (i.e., a regular tf.Variable).


def is_leap_year(years):
    years = tf.convert_to_tensor(years, tf.int32)

    def divides_by(n):
        return tf.math.equal(years % n, 0)
    return tf.math.logical_and(divides_by(4), tf.math.logical_or(~divides_by(100), divides_by(400)))

Calculates whether years are leap years.

Args:
        years: Tensor of int32 type. Elements should be positive.

Returns:
        Tensor of bool type.


def _is_instance(self, obj, class_spec):
    cls = obj.cls
    if isinstance(obj, abstract.AMBIGUOUS_OR_EMPTY) or isinstance(cls, abstract.AMBIGUOUS_OR_EMPTY):
        return None
    return abstract_utils.check_against_mro(self.ctx, cls, class_spec)

Check if the object matches a class specification.

Args:
        obj: A BaseValue, generally the left hand side of an isinstance() call.
        class_spec: A BaseValue, generally the right hand side of an isinstance()
        call.

Returns:
        True if the object is derived from a class in the class_spec, False if
        it is not, and None if it is ambiguous whether obj matches class_spec.


def reduce_min(values, index, name='segmented_reduce_min'):
    return _segment_reduce(values, index, 'amin', name)

Computes the minimum over segments.

    This operations computes the minimum over segments, with support for:

        - Batching using the first dimensions [B1, B2, ..., Bn]. Each element in a batch can have different indices.
        - Vectorization using the last dimension [V1, V2, ...]. If they are present, the output will be an element-wise
        minimum of vectors rather than scalars.

    Only the middle dimensions [I1, ..., Ik] are reduced by the operation.

Args:
        values (`torch.Tensor` of shape [B1, B2, ..., Bn, I1, .., Ik, V1, V2, ..]):
            Tensor containing the values of which the min must be taken segment-wise.
        index (`IndexMap`, indices are of shape [B1, B2, ..., Bn, I1, .., Ik].):
            Index defining the segments.
        name (`str`, *optional*, defaults to 'segmented_reduce_sum'):
            Name for the operation. Currently not used

Returns:
        output_values (`torch.Tensor`of shape [B1, B2, ..., Bn, num_segments, V1, V2, ..]): Tensor containing the
        output values. output_index (`IndexMap`): IndexMap with shape [B1, B2, ..., Bn, num_segments].


def get_special_tokens_mask(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None, already_has_special_tokens: bool=False) -> List[int]:
    if already_has_special_tokens:
        return super().get_special_tokens_mask(token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True)
    if not self.add_bos_token:
        return super().get_special_tokens_mask(token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=False)
    if token_ids_1 is None:
        return [1] + [0] * len(token_ids_0)
    return [1] + [0] * len(token_ids_0) + [1] + [0] * len(token_ids_1)

Retrieves sequence ids from a token list that has no special tokens added. This method is called when adding
    special tokens using the tokenizer `prepare_for_model` or `encode_plus` methods.

Args:
        token_ids_0 (`List[int]`):
            List of IDs.
        token_ids_1 (`List[int]`, *optional*):
            Optional second list of IDs for sequence pairs.
        already_has_special_tokens (`bool`, *optional*, defaults to `False`):
            Whether or not the token list is already formatted with special tokens for the model.

Returns:
        `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.


def resize_image_with_pad_v2(image, target_height, target_width, method=ResizeMethod.BILINEAR, antialias=False):

    def _resize_fn(im, new_size):
        return resize_images_v2(im, new_size, method, antialias=antialias)
    return _resize_image_with_pad_common(image, target_height, target_width, _resize_fn)

Resizes and pads an image to a target width and height.

    Resizes an image to a target width and height by keeping
    the aspect ratio the same without distortion. If the target
    dimensions don't match the image dimensions, the image
    is resized and then padded with zeroes to match requested
    dimensions.

Args:
        image: 4-D Tensor of shape `[batch, height, width, channels]` or 3-D Tensor
        of shape `[height, width, channels]`.
        target_height: Target height.
        target_width: Target width.
        method: Method to use for resizing image. See `image.resize()`
        antialias: Whether to use anti-aliasing when resizing. See 'image.resize()'.

Raises:
        ValueError: if `target_height` or `target_width` are zero or negative.

Returns:
        Resized and padded image.
        If `images` was 4-D, a 4-D float Tensor of shape
        `[batch, new_height, new_width, channels]`.
        If `images` was 3-D, a 3-D float Tensor of shape
        `[new_height, new_width, channels]`.