def _MeanAggregator(inputs, segments):
    result = []
    for inputs_i, segments_i in zip(array_ops.split(inputs, inputs.shape[0]), array_ops.split(segments, segments.shape[0])):
        means_i = math_ops.unsorted_segment_mean(inputs_i, segments_i, num_segments=math_ops.reduce_max(segments_i) + 1)
        result.append(array_ops.reshape(array_ops.gather(means_i, segments_i), [-1]))
    return array_ops_stack.stack(result, axis=0)

Replaces each segment with its mean along the last axis.

    Specifically, each value in the `inputs` tensor gets replaced by the mean
    value computed from the values that belong to the same segment.

Args:
    inputs: A 2-tensor. Aggregation is done over dimension 1.
    segments: A 2-tensor, same shape as `input`.

Returns:
        The result, same shape and type as `inputs`.


def check_copies(overwrite: bool=False, file: Optional[str]=None):
    buffer = {}
    if file is None:
        all_files = glob.glob(os.path.join(TRANSFORMERS_PATH, '**/*.py'), recursive=True)
        all_test_files = glob.glob(os.path.join(MODEL_TEST_PATH, '**/*.py'), recursive=True)
        all_files = list(all_files) + list(all_test_files)
    else:
        all_files = [file]
    diffs = []
    for filename in all_files:
        new_diffs = is_copy_consistent(filename, overwrite, buffer)
        diffs += [f'- {filename}: copy does not match {d[0]} at line {d[1]}' for d in new_diffs]
    if not overwrite and len(diffs) > 0:
        diff = '\n'.join(diffs)
        raise Exception('Found the following copy inconsistencies:\n' + diff + '\nRun `make fix-copies` or `python utils/check_copies.py --fix_and_overwrite` to fix them.')

Check every file is copy-consistent with the original. Also check the model list in the main README and other
    READMEs are consistent.

Args:
        overwrite (`bool`, *optional*, defaults to `False`):
            Whether or not to overwrite the copies when they don't match.
        file (`bool`, *optional*):
            The path to a specific file to check and/or fix.


def pixel_shuffle(self, vision_features: torch.Tensor, scale_factor: float=0.5):
    batch_size, width, height, channels = vision_features.size()
    if height % scale_factor != 0 or width % scale_factor != 0:
        raise ValueError('Height and width must be divisible by scale_factor for proper downsampling.')
    vision_features = vision_features.view(batch_size, width, int(height * scale_factor), int(channels / scale_factor))
    vision_features = vision_features.permute(0, 2, 1, 3).contiguous()
    vision_features = vision_features.view(batch_size, int(height * scale_factor), int(width * scale_factor), int(channels / scale_factor ** 2))
    vision_features = vision_features.permute(0, 2, 1, 3).contiguous()
    return vision_features

Perform pixel shuffle downsampling on vision features.

Args:
        vision_features (`torch.Tensor`):
            Input tensor of shape (batch_size, width, height, channels).
        scale_factor (`float`, *optional*, defaults to `0.5`):
            Factor by which to downsample. Default is 0.5, which halves the dimensions.

Returns:
        vision_features (`torch.Tensor`):
            Downsampled tensor of shape (batch_size, height*scale_factor, width*scale_factor, channels/(scale_factor^2)).


def combine(**kwargs):
    if not kwargs:
        return [OrderedDict()]
    sort_by_key = lambda k: k[0]
    kwargs = OrderedDict(sorted(kwargs.items(), key=sort_by_key))
    first = list(kwargs.items())[0]
    rest = dict(list(kwargs.items())[1:])
    rest_combined = combine(**rest)
    key = first[0]
    values = first[1]
    if not isinstance(values, list):
        values = [values]
    return [OrderedDict(sorted(list(combined.items()) + [(key, v)], key=sort_by_key)) for v in values for combined in rest_combined]

Generate combinations based on its keyword arguments.

    Two sets of returned combinations can be concatenated using +.  Their product
    can be computed using `times()`.

Args:
        **kwargs: keyword arguments of form `option=[possibilities, ...]`
        or `option=the_only_possibility`.

Returns:
        a list of dictionaries for each combination. Keys in the dictionaries are
        the keyword argument names.  Each key has one value - one of the
        corresponding keyword argument values.


def correct_pad(inputs, kernel_size):
    img_dim = 2 if backend.image_data_format() == 'channels_first' else 1
    input_size = inputs.shape[img_dim:img_dim + 2]
    if isinstance(kernel_size, int):
        kernel_size = (kernel_size, kernel_size)
    if input_size[0] is None:
        adjust = (1, 1)
    else:
        adjust = (1 - input_size[0] % 2, 1 - input_size[1] % 2)
    correct = (kernel_size[0] // 2, kernel_size[1] // 2)
    return ((correct[0] - adjust[0], correct[0]), (correct[1] - adjust[1], correct[1]))

Returns a tuple for zero-padding for 2D convolution with downsampling.

Args:
        inputs: Input tensor.
        kernel_size: An integer or tuple/list of 2 integers.

Returns:
        A tuple.


def split_image(self, image: np.ndarray, input_data_format: Optional[Union[str, ChannelDimension]]=None):
    height, width = get_image_size(image, input_data_format)
    mid_width = width // 2
    mid_height = height // 2
    return [self._crop(image, 0, 0, mid_width, mid_height, input_data_format), self._crop(image, mid_width, 0, width, mid_height, input_data_format), self._crop(image, 0, mid_height, mid_width, height, input_data_format), self._crop(image, mid_width, mid_height, width, height, input_data_format), image]

Split an image into 4 equal sub-images, and the concatenate that sequence with the original image.
    That means that a single image becomes a sequence of 5 images.
    This is a "trick" to spend more compute on each image with no changes in the vision encoder.

Args:
        image (`np.ndarray`):
            Images to split.
        input_data_format (`ChannelDimension` or `str`, *optional*):
            The channel dimension format of the input image. If not provided, it will be inferred.


def _calculate_aggregation_loss_known(logits_aggregation, aggregate_mask, aggregation_labels, use_answer_as_supervision, num_aggregation_labels):
    if use_answer_as_supervision:
        target_aggregation = tf.zeros_like(aggregate_mask, dtype=tf.int32)
    else:
        target_aggregation = aggregation_labels
    one_hot_labels = tf.one_hot(target_aggregation, depth=num_aggregation_labels, dtype=tf.float32)
    log_probs = tf.nn.log_softmax(logits_aggregation, axis=-1)
    per_example_aggregation_intermediate = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1)
    if use_answer_as_supervision:
        return per_example_aggregation_intermediate * (1 - aggregate_mask)
    else:
        return per_example_aggregation_intermediate

Calculates aggregation loss when its type is known during training.

    In the weakly supervised setting, the only known information is that for cell selection examples, "no aggregation"
    should be predicted. For other examples (those that require aggregation), no loss is accumulated. In the setting
    where aggregation type is always known, standard cross entropy loss is accumulated for all examples

Args:
        logits_aggregation (`tf.Tensor` of shape `(batch_size, num_aggregation_labels)`):
            Logits per aggregation operation.
        aggregate_mask (`tf.Tensor` of shape `(batch_size, )`):
            A mask set to 1 for examples that should use aggregation functions.
        aggregation_labels (`tf.Tensor` of shape `(batch_size, )`):
            Aggregation function id for every example in the batch.
        use_answer_as_supervision (`bool`, *optional*):
            Whether to use the answer as the only supervision for aggregation examples.
        num_aggregation_labels (`int`, *optional*, defaults to 0):
            The number of aggregation operators to predict.

Returns:
        aggregation_loss_known (`tf.Tensor` of shape `(batch_size,)`): Aggregation loss (when its type is known during
        training) per example.


def __init__(self, input_bytes):
    self.fdp = atheris.FuzzedDataProvider(input_bytes)

FuzzingHelper initializer.

Args:
        input_bytes: Input randomized bytes used to create a FuzzedDataProvider.


def scatter_div(self, sparse_delta, use_locking=False, name=None):
    if not isinstance(sparse_delta, indexed_slices.IndexedSlices):
        raise TypeError('sparse_delta is not IndexedSlices: %s' % sparse_delta)
    return gen_state_ops.scatter_div(self._variable, sparse_delta.indices, sparse_delta.values, use_locking=use_locking, name=name)

Divide this variable by `tf.IndexedSlices`.

Args:
        sparse_delta: `tf.IndexedSlices` to divide this variable by.
        use_locking: If `True`, use locking during the operation.
        name: the name of the operation.

Returns:
        A `Tensor` that will hold the new value of this variable after
        the scattered division has completed.

Raises:
        TypeError: if `sparse_delta` is not an `IndexedSlices`.


def interp(x: Array['*d'], from_: Tuple[_MinMaxValue, _MinMaxValue], to: Tuple[_MinMaxValue, _MinMaxValue], *, axis: int=-1, xnp: numpy_utils.NpModule=...) -> FloatArray['*d']:
    if axis != -1:
        raise NotImplementedError('Only last axis supported for now. Please send a feature request.')
    from_ = tuple((xnp.asarray(v) for v in from_))
    to = tuple((xnp.asarray(v) for v in to))
    a, b = _linear_interp_factors(*from_, *to)
    return a * x + b

Linearly scale the given value by the given range.

    Somehow similar to `np.interp` or `scipy.interpolate.inter1d` with some
    differences like support scaling an axis by a different factors and
    extrapolate values outside the boundaries.

    `from_` and `to` are expected to be `(min, max)` tuples and the function
    interpolate between the two ranges.

    Example: Normalizing a uint8 image to `(-1, 1)`.

    ```python
    img = jnp.array([
        [0, 0],
        [127, 255],
    ])
    img = enp.interp(img, (0, 255), (0, 1))
    img == jnp.array([
        [-1, -1],
        [0.498..., 1],
    ])
    ```

    `min` and `max` can be either float values or array like structure, in which
    case the numpy broadcasting rules applies (x should be a `Array[... d]` and
    min/max values should be `Array[d]`.

    Example: Converting normalized 3d coordinates to world coordinates.

    ```python
    coords = enp.interp(coords, from_=(-1, 1), to=(0, (h, w, d)))
    ```

    * `coords[:, 0]` is interpolated from `(-1, 1)` to `(0, h)`
    * `coords[:, 1]` is interpolated from `(-1, 1)` to `(0, w)`
    * `coords[:, 2]` is interpolated from `(-1, 1)` to `(0, d)`

Args:
        x: Array to scale
        from_: Range of x.
        to: Range to which normalize x.
        axis: Axis on which normalizing. Only relevant if `from_` or `to` items
        contains range value.
        xnp: Numpy module to use

Returns:
        Float tensor with same shape as x, but with normalized coordinates.


def forward(self, pixel_values: Optional[torch.Tensor]=None, noise: Optional[torch.FloatTensor]=None, head_mask: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, interpolate_pos_encoding: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[tuple, HieraForPreTrainingOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    outputs = self.hiera(pixel_values, noise=noise, head_mask=head_mask, output_attentions=output_attentions, output_hidden_states=True, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=return_dict)
    feature_maps = outputs[-1]
    bool_masked_pos = outputs[1]
    ids_to_restore = outputs[2]
    feature_maps = feature_maps[1:self.hiera.config.num_query_pool + 1] + (feature_maps[-1],)
    fused_hidden_states = self.multiscale_fusion(feature_maps)
    fused_hidden_states = self.encoder_norm(fused_hidden_states)
    logits, bool_masked_pos = self.decoder(fused_hidden_states, bool_masked_pos=bool_masked_pos, head_mask=head_mask, output_attentions=output_attentions)
    loss = self.forward_loss(pixel_values, logits, bool_masked_pos)
    if not return_dict:
        output = (logits, bool_masked_pos, ids_to_restore)
        if output_hidden_states:
            output = output + (outputs[3],)
        if output_attentions:
            output = output + (outputs[4],)
        if output_hidden_states:
            output = output + (outputs[-1],)
        return (loss,) + output if loss is not None else output
    return HieraForPreTrainingOutput(loss=loss, logits=logits, bool_masked_pos=bool_masked_pos, ids_restore=ids_to_restore, hidden_states=outputs.hidden_states if output_hidden_states else None, attentions=outputs.attentions, reshaped_hidden_states=outputs.reshaped_hidden_states if output_hidden_states else None)

noise (`torch.FloatTensor` of shape `(batch_size, num_mask_units)`, *optional*):
    Mainly used for testing purposes to control randomness and maintain the reproducibility

Example:
    ```python
    >>> from transformers import AutoImageProcessor, HieraForPreTraining
    >>> import torch
    >>> from PIL import Image
    >>> import requests

    >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    >>> image = Image.open(requests.get(url, stream=True).raw)

    >>> image_processor = AutoImageProcessor.from_pretrained("facebook/hiera-tiny-224-mae-hf")
    >>> model = HieraForPreTraining.from_pretrained("facebook/hiera-tiny-224-mae-hf")

    >>> inputs = image_processor(images=image, return_tensors="pt")

    >>> outputs = model(**inputs)
    >>> logits = outputs.logits
    >>> loss = outputs.loss
    >>> print(list(logits.shape))
    [1, 196, 768]
    ```


def round(x):
    return math_ops.round(x)

Element-wise rounding to the closest integer.

    In case of tie, the rounding mode used is "half to even".

Args:
        x: Tensor or variable.

Returns:
        A tensor.


def gradient_override_map(self, op_type_map) -> Iterator[None]:
    if not isinstance(op_type_map, dict):
        raise TypeError('op_type_map must be a dictionary mapping strings to strings')
    saved_mappings = {}
    for op_type, mapped_op_type in op_type_map.items():
        if not (isinstance(op_type, str) and isinstance(mapped_op_type, str)):
            raise TypeError('op_type_map must be a dictionary mapping strings to strings')
        try:
            saved_mappings[op_type] = self._gradient_override_map[op_type]
        except KeyError:
            pass
        self._gradient_override_map[op_type] = mapped_op_type
    try:
        yield
    finally:
        for op_type, mapped_op_type in op_type_map.items():
            try:
                self._gradient_override_map[op_type] = saved_mappings[op_type]
            except KeyError:
                del self._gradient_override_map[op_type]

EXPERIMENTAL: A context manager for overriding gradient functions.

    This context manager can be used to override the gradient function
    that will be used for ops within the scope of the context.

    For example:

    ```python
    @tf.RegisterGradient("CustomSquare")
    def _custom_square_grad(op, grad):
        # ...

    with tf.Graph().as_default() as g:
        c = tf.constant(5.0)
        s_1 = tf.square(c)  # Uses the default gradient for tf.square.
        with g.gradient_override_map({"Square": "CustomSquare"}):
        s_2 = tf.square(s_2)  # Uses _custom_square_grad to compute the
                            # gradient of s_2.
    ```

Args:
        op_type_map: A dictionary mapping op type strings to alternative op type
        strings.

Returns:
        A context manager that sets the alternative op type to be used for one
        or more ops created in that context.

Raises:
        TypeError: If `op_type_map` is not a dictionary mapping strings to
        strings.


def __getitem__(self, key):
    if key is None:
        key = self._key()
    value = self._get_recursive(key)
    if value is None:
        value = self[key] = self.default_factory()
    return value

Gets the value at key (or current context), or sets default value.

Args:
        key: May be `None` or `Graph`object. When `None`, the key is set to the
        current context.

Returns:
        Either the cached or default value.


def __init__(self, config: Owlv2Config):
    super().__init__()
    self.layers = nn.ModuleList([Owlv2EncoderLayer(config) for _ in range(config.num_hidden_layers)])
    self.gradient_checkpointing = False

Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
    [`Owlv2EncoderLayer`].

Args:
        config: Owlv2Config


def bessel_j0(x, name=None):
    with ops.name_scope(name, 'bessel_j0', [x]):
        return gen_special_math_ops.bessel_j0(x)

Computes the Bessel j0 function of `x` element-wise.

    Modified Bessel function of order 0.

    >>> tf.math.special.bessel_j0([0.5, 1., 2., 4.]).numpy()
    array([ 0.93846981,  0.76519769,  0.22389078, -0.39714981], dtype=float32)

Args:
        x: A `Tensor` or `SparseTensor`. Must be one of the following types: `half`,
        `float32`, `float64`.
        name: A name for the operation (optional).

Returns:
        A `Tensor` or `SparseTensor`, respectively. Has the same type as `x`.

    @compatibility(scipy)
    Equivalent to scipy.special.j0
    @end_compatibility


def get_plugin_asset(plugin_asset_cls, graph=None):
    if graph is None:
        graph = ops.get_default_graph()
    if not plugin_asset_cls.plugin_name:
        raise ValueError('Class %s has no plugin_name' % plugin_asset_cls.__name__)
    name = _PLUGIN_ASSET_PREFIX + plugin_asset_cls.plugin_name
    container = graph.get_collection(name)
    if container:
        if len(container) != 1:
            raise ValueError('Collection for %s had %d items, expected 1' % (name, len(container)))
        instance = container[0]
        if not isinstance(instance, plugin_asset_cls):
            raise ValueError('Plugin name collision between classes %s and %s' % (plugin_asset_cls.__name__, instance.__class__.__name__))
    else:
        instance = plugin_asset_cls()
        graph.add_to_collection(name, instance)
        graph.add_to_collection(_PLUGIN_ASSET_PREFIX, plugin_asset_cls.plugin_name)
    return instance

Acquire singleton PluginAsset instance from a graph.

    PluginAssets are always singletons, and are stored in tf Graph collections.
    This way, they can be defined anywhere the graph is being constructed, and
    if the same plugin is configured at many different points, the user can always
    modify the same instance.

Args:
        plugin_asset_cls: The PluginAsset class
        graph: (optional) The graph to retrieve the instance from. If not specified,
        the default graph is used.

Returns:
        An instance of the plugin_asset_class

Raises:
        ValueError: If we have a plugin name collision, or if we unexpectedly find
        the wrong number of items in a collection.


def scatter_nd_update(self, indices, updates, name=None):
    return self._lazy_read(gen_state_ops.resource_scatter_nd_update(self.handle, indices, ops.convert_to_tensor(updates, self.dtype), name=name))

Applies sparse assignment to individual values or slices in a Variable.

    `ref` is a `Tensor` with rank `P` and `indices` is a `Tensor` of rank `Q`.

    `indices` must be integer tensor, containing indices into `ref`.
    It must be shape `[d_0, ..., d_{Q-2}, K]` where `0 < K <= P`.

    The innermost dimension of `indices` (with length `K`) corresponds to
    indices into elements (if `K = P`) or slices (if `K < P`) along the `K`th
    dimension of `ref`.

    `updates` is `Tensor` of rank `Q-1+P-K` with shape:

    ```
    [d_0, ..., d_{Q-2}, ref.shape[K], ..., ref.shape[P-1]].
    ```

    For example, say we want to add 4 scattered elements to a rank-1 tensor to
    8 elements. In Python, that update would look like this:

    ```python
        ref = tf.Variable([1, 2, 3, 4, 5, 6, 7, 8])
        indices = tf.constant([[4], [3], [1] ,[7]])
        updates = tf.constant([9, 10, 11, 12])
        op = ref.scatter_nd_update(indices, updates)
        with tf.compat.v1.Session() as sess:
        print sess.run(op)
    ```

    The resulting update to ref would look like this:

        [1, 11, 3, 10, 9, 6, 7, 12]

    See `tf.scatter_nd` for more details about how to make updates to
    slices.

Args:
        indices: The indices to be used in the operation.
        updates: The values to be used in the operation.
        name: the name of the operation.

Returns:
        The updated variable.


def num_accelerators(self, task_type=None, task_id=None, config_proto=None):
    master = self.master(task_type, task_id)
    devices = get_accelerator_devices(master, config_proto)
    mapping = collections.defaultdict(int)
    for device in devices:
        if task_type is not None and task_id is not None:
            job_path = '/job:%s' % task_type
            task_path = '/task:%s' % task_id
            if job_path not in device.name or task_path not in device.name:
                continue
        mapping[device.device_type] += 1
    return mapping

Returns the number of accelerator cores per worker.

    This returns the number of accelerator cores (such as GPUs and TPUs)
    available per worker.

    Optionally, we allow callers to specify the task_type, and task_id, for
    if they want to target a specific TensorFlow task to query
    the number of accelerators. This is to support heterogenous environments,
    where the number of accelerators cores per host is different.

Args:
        task_type: (Optional) The type of the TensorFlow task of the machine we
        want to query.
        task_id: (Optional) The index of the TensorFlow task of the machine we
        want to query.
        config_proto: (Optional) Configuration for starting a new session to
        query how many accelerator cores it has.

Returns:
        A map of accelerator types to number of cores.


def compute_mask(self, inputs, mask=None):
    if not self._supports_masking:
        if any((m is not None for m in nest.flatten(mask))):
            raise TypeError('Layer ' + self.name + ' does not support masking, but was passed an input_mask: ' + str(mask))
        return None
    return mask

Computes an output mask tensor.

Args:
        inputs: Tensor or list of tensors.
        mask: Tensor or list of tensors.

Returns:
        None or a tensor (or list of tensors,
            one per output tensor of the layer).


def fit(weights: Array, train_dataset: Dataset, iters: int, learning_rate: float, log_span: int, val_dataset: typing.Optional[Dataset]=None) -> Array:
    grad_loss = jit(grad(cross_entropy_loss, argnums=0))
    for t in range(iters):
        weights = weights - learning_rate * grad_loss(weights, train_dataset.X, train_dataset.Y)
        if (t + 1) % log_span != 0:
            continue
        metrics_train = jit(get_metrics)(weights, train_dataset)
        print()
        print('iter:\t%d' % (t + 1))
        print()
        print('train accuracy:\t%.5f' % metrics_train.accuracy)
        print('train prec.:\t%.5f' % metrics_train.precision)
        print('train recall:\t%.5f' % metrics_train.recall)
        print('train fscore:\t%.5f' % metrics_train.fscore)
        print('train loss:\t%.5f' % metrics_train.loss)
        print()
        if val_dataset is None:
            continue
        metrics_val = jit(get_metrics)(weights, val_dataset)
        print('val accuracy:\t%.5f' % metrics_val.accuracy)
        print('val prec.:\t%.5f' % metrics_val.precision)
        print('val recall:\t%.5f' % metrics_val.recall)
        print('val fscore:\t%.5f' % metrics_val.fscore)
        print('val loss:\t%.5f' % metrics_val.loss)
        print()
    return weights

Updates the weights with the given dataset.

Args:
        weights: A weight vector.
        train_dataset: A train dataset.
        iters: A number of iterations.
        learning_rate: A learning rate.
        log_span: A span to log metrics.
        val_dataset: A validation dataset (optional).

Returns:
        An updated weight vector.


def __init__(self, padding=1, **kwargs):
    super(ZeroPadding1D, self).__init__(**kwargs)
    self.padding = conv_utils.normalize_tuple(padding, 2, 'padding')
    self.input_spec = InputSpec(ndim=3)

Zero-padding layer for 1D input (e.g. temporal sequence).

Example:
    >>> input_shape = (2, 2, 3)
    >>> x = np.arange(np.prod(input_shape)).reshape(input_shape)
    >>> print(x)
    [[[ 0  1  2]
        [ 3  4  5]]
    [[ 6  7  8]
        [ 9 10 11]]]
    >>> y = tf.keras.layers.ZeroPadding1D(padding=2)(x)
    >>> print(y)
    tf.Tensor(
        [[[ 0  0  0]
        [ 0  0  0]
        [ 0  1  2]
        [ 3  4  5]
        [ 0  0  0]
        [ 0  0  0]]
        [[ 0  0  0]
        [ 0  0  0]
        [ 6  7  8]
        [ 9 10 11]
        [ 0  0  0]
        [ 0  0  0]]], shape=(2, 6, 3), dtype=int64)

Args:
        padding: Int, or tuple of int (length 2), or dictionary.
            - If int:
            How many zeros to add at the beginning and end of
            the padding dimension (axis 1).
            - If tuple of int (length 2):
            How many zeros to add at the beginning and the end of
            the padding dimension (`(left_pad, right_pad)`).

    Input shape:
        3D tensor with shape `(batch_size, axis_to_pad, features)`

    Output shape:
        3D tensor with shape `(batch_size, padded_axis, features)`


def _num_present(losses, weights, per_batch=False):
    if isinstance(weights, float) and weights != 0.0 or (context.executing_eagerly() and weights._rank() == 0 and (not math_ops.equal(weights, 0.0))):
        return _num_elements(losses)
    with ops.name_scope(None, 'num_present', (losses, weights)) as scope:
        weights = math_ops.cast(weights, dtype=dtypes.float32)
        present = array_ops.where(math_ops.equal(weights, 0.0), array_ops.zeros_like(weights), array_ops.ones_like(weights))
        present = weights_broadcast_ops.broadcast_weights(present, losses)
        if per_batch:
            return math_ops.reduce_sum(present, axis=math_ops.range(1, array_ops.rank(present)), keepdims=True, name=scope)
        return math_ops.reduce_sum(present, name=scope)

Computes the number of elements in the loss function induced by `weights`.

    A given weights tensor induces different numbers of usable elements in the
    `losses` tensor. The `weights` tensor is broadcast across `losses` for all
    possible dimensions. For example, if `losses` is a tensor of dimension
    `[4, 5, 6, 3]` and `weights` is a tensor of shape `[4, 5]`, then `weights` is,
    in effect, tiled to match the shape of `losses`. Following this effective
    tile, the total number of present elements is the number of non-zero weights.

Args:
        losses: `Tensor` of shape `[batch_size, d1, ... dN]`.
        weights: `Tensor` of shape `[]`, `[batch_size]` or
        `[batch_size, d1, ... dK]`, where K < N.
        per_batch: Whether to return the number of elements per batch or as a sum
        total.

Returns:
        The number of present (non-zero) elements in the losses tensor. If
        `per_batch` is `True`, the value is returned as a tensor of size
        `[batch_size]`. Otherwise, a single scalar tensor is returned.


def list_key_values(input: t.Dict[str, str]) -> None:
    for cmd, desc in input.items():
        print(f'{cmd}  =>  {desc}')

Display key-value pairs from a dictionary.

Args:
        input (Dict[str, str]): The dictionary containing key-value pairs.


def _post_process_masks_pt(self, masks, original_sizes, reshaped_input_sizes, mask_threshold=0.0, binarize=True, pad_size=None):
    requires_backends(self, ['torch'])
    pad_size = self.pad_size if pad_size is None else pad_size
    target_image_size = (pad_size['height'], pad_size['width'])
    if isinstance(original_sizes, (torch.Tensor, np.ndarray)):
        original_sizes = original_sizes.tolist()
    if isinstance(reshaped_input_sizes, (torch.Tensor, np.ndarray)):
        reshaped_input_sizes = reshaped_input_sizes.tolist()
    output_masks = []
    for i, original_size in enumerate(original_sizes):
        if isinstance(masks[i], np.ndarray):
            masks[i] = torch.from_numpy(masks[i])
        elif not isinstance(masks[i], torch.Tensor):
            raise ValueError('Input masks should be a list of `torch.tensors` or a list of `np.ndarray`')
        interpolated_mask = F.interpolate(masks[i], target_image_size, mode='bilinear', align_corners=False)
        interpolated_mask = interpolated_mask[..., :reshaped_input_sizes[i][0], :reshaped_input_sizes[i][1]]
        interpolated_mask = F.interpolate(interpolated_mask, original_size, mode='bilinear', align_corners=False)
        if binarize:
            interpolated_mask = interpolated_mask > mask_threshold
        output_masks.append(interpolated_mask)
    return output_masks

Remove padding and upscale masks to the original image size.

Args:
        masks (`Union[List[torch.Tensor], List[np.ndarray]]`):
            Batched masks from the mask_decoder in (batch_size, num_channels, height, width) format.
        original_sizes (`Union[torch.Tensor, List[Tuple[int,int]]]`):
            The original sizes of each image before it was resized to the model's expected input shape, in (height,
            width) format.
        reshaped_input_sizes (`Union[torch.Tensor, List[Tuple[int,int]]]`):
            The size of each image as it is fed to the model, in (height, width) format. Used to remove padding.
        mask_threshold (`float`, *optional*, defaults to 0.0):
            The threshold to use for binarizing the masks.
        binarize (`bool`, *optional*, defaults to `True`):
            Whether to binarize the masks.
        pad_size (`int`, *optional*, defaults to `self.pad_size`):
            The target size the images were padded to before being passed to the model. If None, the target size is
            assumed to be the processor's `pad_size`.

Returns:
        (`torch.Tensor`): Batched masks in batch_size, num_channels, height, width) format, where (height, width)
        is given by original_size.


def _compare_constant_tuple_prefix(op, prefix, constant_tuple, reverse):
    length = min(len(prefix), len(constant_tuple))
    trimmed_prefix = prefix[:length]
    trimmed_constant_tuple = constant_tuple[:length]
    if trimmed_prefix == trimmed_constant_tuple:
        if len(prefix) >= len(constant_tuple):
            if reverse:
                return op in (slots.LT, slots.LE, slots.NE)
            else:
                return op in (slots.NE, slots.GE, slots.GT)
        return None
    if reverse:
        return _compare_constants(op, trimmed_constant_tuple, trimmed_prefix)
    else:
        return _compare_constants(op, trimmed_prefix, trimmed_constant_tuple)

Compares a tuple's constant prefix against a constant tuple.

Args:
        op: A comparison operator, such as LT (less than).
        prefix: A constant prefix of a non-constant tuple (referred to as "left" in
        the inline comments). So if left=(3, 2, ...), prefix=(3, 2).
        constant_tuple: A constant tuple (referred to as "right").
        reverse: Whether left and right should be reversed for the comparison.

Returns:
        A bool of the comparison result if it can be determined, None otherwise.


def potentially_ragged_concat(tensors):
    if len(tensors) == 1:
        return tensors[0]
    elif isinstance(tensors[0], tf.SparseTensor):
        return tf.sparse.concat(axis=0, sp_inputs=tensors)
    elif isinstance(tensors[0], tf.RaggedTensor):
        return tf.concat(tensors, axis=0)
    non_batch_shapes = tf.stack([tf.shape(tensor)[1:] for tensor in tensors])
    constant_dims = tf.math.reduce_all(non_batch_shapes == non_batch_shapes[:1], axis=0)
    if tf.math.reduce_all(constant_dims).numpy().item():
        if _is_scalar(tensors[0]):
            return tf.stack(tensors, axis=0)
        else:
            return tf.concat(tensors, axis=0)
    constant_inner_dimensions = constant_dims.numpy().tolist()[::-1].index(False)
    if constant_inner_dimensions == 0:
        constant_inner_shape = None
    else:
        constant_inner_shape = tensors[0].shape[-constant_inner_dimensions:]
    return tf.ragged.constant([tensor.numpy() for tensor in tensors], inner_shape=constant_inner_shape).merge_dims(0, 1)

Concats `Tensor`s along their first dimension.

Args:
        tensors: List of `Tensor`s.

Returns:
        Concatenation of the inputs along the first dimension -- of type
        `np.ndarray` if all input shapes are compatible, or `tf.RaggedTensor`
        if not.


def SmartBroadcastGradientArgs(x, y, grad=None):
    del grad
    x_shape = array_ops.shape(x)
    y_shape = array_ops.shape(y)
    if not context.executing_eagerly() and isinstance(x, tensor.Tensor) and isinstance(y, tensor.Tensor):
        x_axes, y_axes = _InferGradientReductionAxes(x.shape, y.shape)
    else:
        x_axes, y_axes = (None, None)
    if x_axes is None or y_axes is None:
        x_axes, y_axes = gen_array_ops.broadcast_gradient_args(x_shape, y_shape)
        x_must_reduce = True
        y_must_reduce = True
    else:
        x_must_reduce = x_axes or x.shape.rank < y.shape.rank
        y_must_reduce = y_axes or y.shape.rank < x.shape.rank
    return ((x_shape, x_axes, x_must_reduce), (y_shape, y_axes, y_must_reduce))

Version of `BroadcastGradientArgs` optimized for partially-known shapes.

Args:
        x: The first argument of a broadcasting binary op.
        y: The second argument of a broadcasting binary op.
        grad: Deprecated.

Returns:
        A pair of triples, one per argument with
        * Shape of the argument (tensor);
        * Reduction axes for the argument (list or tensor);
        * Boolean indicating whether the reduction must be applied.


def add_inputs(self, mutable_accumulator, elements, *args, **kwargs):
    for element in elements:
        mutable_accumulator = self.add_input(mutable_accumulator, element, *args, **kwargs)
    return mutable_accumulator

Returns the result of folding each element in elements into accumulator.

    This is provided in case the implementation affords more efficient
    bulk addition of elements. The default implementation simply loops
    over the inputs invoking add_input for each one.

Args:
        mutable_accumulator: the current accumulator,
        may be modified and returned for efficiency
        elements: the elements to add, should not be mutated
        *args: Additional arguments and side inputs.
        **kwargs: Additional arguments and side inputs.


def solve(a, b):
    if any_symbolic_tensors((a, b)):
        return Solve().symbolic_call(a, b)
    return _solve(a, b)

Solves a linear system of equations given by `a x = b`.

Args:
        a: A tensor of shape `(..., M, M)` representing the coefficients matrix.
        b: A tensor of shape `(..., M)` or `(..., M, N)` representing the
        right-hand side or "dependent variable" matrix.

Returns:
        A tensor of shape `(..., M)` or `(..., M, N)` representing the solution
        of the linear system. Returned shape is identical to `b`.


def resume(resume_delay=0):
    return ProcessContinuation(resume_delay=resume_delay)

A convenient method that produces a ``ProcessContinuation``.

Args:
        resume_delay: delay after which processing current element should be
        resumed.
    Returns: a ``ProcessContinuation`` for signalling the runner that current
        input element has not been fully processed and should be resumed later.


def enable_eager_execution_internal(config=None, device_policy=None, execution_mode=None, server_def=None) -> None:
    if config is not None and (not isinstance(config, config_pb2.ConfigProto)):
        raise TypeError('config must be a tf.ConfigProto, but got %s' % type(config))
    if device_policy not in (None, context.DEVICE_PLACEMENT_EXPLICIT, context.DEVICE_PLACEMENT_WARN, context.DEVICE_PLACEMENT_SILENT, context.DEVICE_PLACEMENT_SILENT_FOR_INT32):
        raise ValueError('device_policy must be one of None, DEVICE_PLACEMENT_*')
    if execution_mode not in (None, context.SYNC, context.ASYNC):
        raise ValueError('execution_mode must be one of None, SYNC, ASYNC')
    if context.default_execution_mode == context.GRAPH_MODE:
        graph_mode_has_been_used = _default_graph_stack._global_default_graph is not None
        if graph_mode_has_been_used:
            raise ValueError('tf.enable_eager_execution must be called at program startup.')
    context.default_execution_mode = context.EAGER_MODE
    with context._context_lock:
        if context._context is None:
            context._set_context_locked(context.Context(config=config, device_policy=device_policy, execution_mode=execution_mode, server_def=server_def))
        elif config is not None and config is not context._context._config or (device_policy is not None and device_policy is not context._context._device_policy) or (execution_mode is not None and execution_mode is not context._context._execution_mode):
            raise ValueError('Trying to change the options of an active eager execution. Context config: %s, specified config: %s. Context device policy: %s, specified device policy: %s. Context execution mode: %s,  specified execution mode %s.' % (context._context._config, config, context._context._device_policy, device_policy, context._context._execution_mode, execution_mode))
        else:
            context._context._thread_local_data.is_eager = True
    context.context = context.context_safe

Enables eager execution for the lifetime of this program.

    Most of the doc string for enable_eager_execution is relevant here as well.

Args:
        config: See enable_eager_execution doc string
        device_policy: See enable_eager_execution doc string
        execution_mode: See enable_eager_execution doc string
        server_def: (Optional.) A tensorflow::ServerDef proto. Enables execution on
        remote devices. GrpcServers need to be started by creating an identical
        server_def to this, and setting the appropriate task_indexes, so that the
        servers can communicate. It will then be possible to execute operations on
        remote devices.

Raises:
        ValueError


def pow(x, y, name=None):
    with ops.name_scope(name, 'Pow', [x]) as name:
        return gen_math_ops._pow(x, y, name=name)

Computes the power of one value to another.

    Given a tensor `x` and a tensor `y`, this operation computes \\(x^y\\) for
    corresponding elements in `x` and `y`. For example:

    ```python
    x = tf.constant([[2, 2], [3, 3]])
    y = tf.constant([[8, 16], [2, 3]])
    tf.pow(x, y)  # [[256, 65536], [9, 27]]
    ```

Args:
        x: A `Tensor` of type `float16`, `float32`, `float64`, `int32`, `int64`,
        `complex64`, or `complex128`.
        y: A `Tensor` of type `float16`, `float32`, `float64`, `int32`, `int64`,
        `complex64`, or `complex128`.
        name: A name for the operation (optional).

Returns:
        A `Tensor`.


def __init__(self, entity=None, key=None):
    self.entity = entity
    self.key = key
    self._pb = FakeMessage(entity, key)

Fake mutation request object.

    Requires exactly one of entity or key to be set.

Args:
        entity: (``google.cloud.datastore.entity.Entity``) entity representing
        this upsert mutation
        key: (``google.cloud.datastore.key.Key``) key representing
        this delete mutation


def convert_pil_frames_to_video(videos: List[VideoInput]) -> List[Union['np.ndarray', 'torch.Tensor']]:
    if not isinstance(videos[0], (list, tuple)):
        return videos
    video_converted = []
    for video in videos:
        video = [np.array(frame) for frame in video]
        video = np.stack(video)
        video_converted.append(video)
    return video_converted

Given a batch of videos, converts each video to a 4D array. If video is already in array type,
    it is simply returned. We assume that all inputs in the list are in the same format, based on the type of the first element.

Args:
        videos (`VideoInput`):
            Video inputs to turn into a list of videos.


def basis(sample_paths):
    samples = tf.convert_to_tensor(sample_paths)
    dim = samples.shape.as_list()[-1]
    grid = tf.range(0, degree + 1, dtype=samples.dtype)
    samples_centered = samples - tf.math.reduce_mean(samples, axis=0)
    samples_centered = tf.expand_dims(samples_centered, -2)
    grid = tf.meshgrid(*dim * [grid])
    grid = tf.reshape(tf.stack(grid, -1), [-1, dim])
    basis_expansion = tf.reduce_prod(samples_centered ** grid, -1)
    return tf.transpose(basis_expansion)

Computes polynomial basis expansion at the given sample points.

Args:
        sample_paths: A `Tensor`s of either `flot32` or `float64` dtype and of
        shape `[num_samples, dim]` where `dim` has to be statically known.

Returns:
        A `Tensor`s of shape `[degree * dim, num_samples]`.


def push(self, is_building_function, enter_context_fn, device_stack):
    self.stack.append(ContextSwitch(is_building_function, enter_context_fn, device_stack))

Push metadata about a context switch onto the stack.

    A context switch can take any one of the two forms: installing a graph as
    the default graph, or entering the eager context. For each context switch,
    we record whether or not the entered context is building a function.

Args:
        is_building_function: (bool.) Whether the context is building a function.
        enter_context_fn: (function.) A callable that executes the context switch.
        For example, `graph.as_default` or `eager_mode`.
        device_stack: If applicable, the device function stack for this graph.
        When breaking out of graphs in init_scope, the innermost nonempty device
        stack is used. Eager contexts put `None` here and the value is never
        used.


def interpolate(hidden_states, ratio):
    batch_size, time_length, classes_num = hidden_states.shape
    upsampled = hidden_states[:, :, None, :].repeat(1, 1, ratio, 1)
    upsampled = upsampled.reshape(batch_size, time_length * ratio, classes_num)
    return upsampled

Interpolate data in time domain. This is used to compensate the resolution reduction in downsampling of a CNN.

Args:
        hidden_states (`torch.FloatTensor` of shape (batch_size, time_length, classes_num)):
            Input hidden states
        ratio (`int`):
            The ratio of the length of the output to the length of the input.


def most_specific_compatible_shape(self, other) -> 'TensorShape':
    other = as_shape(other)
    if self.dims is None or other.dims is None or self.rank != other.rank:
        return unknown_shape()
    dims = [d1 if d1 is not None and d2 is not None and (d1 == d2) else None for d1, d2 in zip(self.dims, other.dims)]
    return TensorShape(dims)

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

    * TensorShape([None, 1]) is the most specific TensorShape compatible with
        both TensorShape([2, 1]) and TensorShape([5, 1]). Note that
        TensorShape(None) is also compatible with above mentioned TensorShapes.

    * TensorShape([1, 2, 3]) is the most specific TensorShape compatible with
        both TensorShape([1, 2, 3]) and TensorShape([1, 2, 3]). There are more
        less specific TensorShapes compatible with above mentioned TensorShapes,
        e.g. TensorShape([1, 2, None]), TensorShape(None).

Args:
        other: Another `TensorShape`.

Returns:
        A `TensorShape` which is the most specific compatible shape of `self`
        and `other`.


def _replacer(self, match: re.Match[str], is_verbatim: bool, is_global: bool) -> str:
    symbol_name = match.group(0)
    if symbol_name in self._local_symbol_replacement_cache:
        return self._local_symbol_replacement_cache[symbol_name]
    if symbol_name in self._global_symbol_replacement_cache:
        return self._global_symbol_replacement_cache[symbol_name]
    if is_verbatim:
        declaration_replacement = symbol_name
        reference_replacement = symbol_name
    else:
        capture_name = self._generate_unique_name(symbol_name)
        capture_pattern = '[^ ]+'
        maybe_global_flag = '$' if is_global else ''
        declaration_replacement = f'[[{maybe_global_flag}{capture_name}:{capture_pattern}]]'
        reference_replacement = f'[[{maybe_global_flag}{capture_name}]]'
    if is_global:
        self._global_symbol_replacement_cache[symbol_name] = reference_replacement
    else:
        self._local_symbol_replacement_cache[symbol_name] = reference_replacement
    return declaration_replacement

A symbol-name replacement function for use in `re.sub`.

Args:
        match: The match object produced by `self._SYMBOL_NAME_REGEX`.
        is_verbatim: Whether the newly matched symbol appears in a "CHECK-LABEL"
        directive, in which case it should be checked verbatim (not replaced
        with a regex capture).
        is_global: Whether the newly matched symbol appears in a declaration at
        global scope, i.e. whether it's a function name. If so, it should be
        remembered across function boundaries.

Returns:
        The replacement string for the symbol name.


def _model_setup():
    context.set_log_device_placement(True)
    batch_size = 64
    steps = 2
    with collective_strategy.CollectiveAllReduceStrategy().scope():
        train_ds, _ = mnist_testing_utils.mnist_synthetic_dataset(batch_size, steps)
        model = mnist_testing_utils.get_mnist_model((28, 28, 1))
    return (batch_size, steps, train_ds, model)

Set up a MNIST Keras model for testing purposes.

    Builds a MNIST Keras model and returns model information.

Returns:
        A tuple of (batch_size, steps, train_dataset, mode)


def __init__(self, name='logcosh', dtype=None):
    super(LogCoshError, self).__init__(logcosh, name, dtype=dtype)

Computes the logarithm of the hyperbolic cosine of the prediction error.

    `logcosh = log((exp(x) + exp(-x))/2)`, where x is the error (y_pred - y_true)

Args:
        name: (Optional) string name of the metric instance.
        dtype: (Optional) data type of the metric result.

    Standalone usage:

    >>> m = tf.keras.metrics.LogCoshError()
    >>> m.update_state([[0, 1], [0, 0]], [[1, 1], [0, 0]])
    >>> m.result().numpy()
    0.10844523

    >>> m.reset_state()
    >>> m.update_state([[0, 1], [0, 0]], [[1, 1], [0, 0]],
    ...                sample_weight=[1, 0])
    >>> m.result().numpy()
    0.21689045

    Usage with `compile()` API:

    ```python
    model.compile(optimizer='sgd',
                loss='mse',
                metrics=[tf.keras.metrics.LogCoshError()])
    ```


def __init__(self, model, preprocess_input=None, input_dtype=None, feature_description=None, **kwargs):
    """
model: model: Base tensorflow model used for TFX-BSL RunInference transform.
preprocess_input: Preprocess method to be included as part of the
  model's serving signature.
input_dtype: tf dtype of the inputs passed to the model.
  For eg: tf.int32, tf.uint8.
feature_description: Feature spec to parse inputs from tf.train.Example
  using tf.parse_example(). For more details, please take a look at
  https://www.tensorflow.org/api_docs/python/tf/io/parse_example
If there are extra arguments(for eg: training=False) that should be
passed to the base tf model during inference, please pass them in kwargs.
"""
    super().__init__()
    self.model = model
    self.preprocess_input = preprocess_input
    self.input_dtype = input_dtype
    self.feature_description = feature_description
    if not feature_description:
        self.feature_description = {'image': tf.io.FixedLenFeature((), tf.string)}
    self._kwargs = kwargs

Helper class used to wrap a based tf.keras.Model object with a serving
    signature that can passed to the tfx_bsl RunInference transform.

    A TF model saved using this helper class expects inputs as
        images serialized to tf.string using tf.io.parse_tensor
        and then passing serialized images to the RunInference transform
        in the tf.train.Example. More about tf.train.Example at
        https://www.tensorflow.org/api_docs/python/tf/train/Example

    Usage:
    Step 1:
    # Save the base TF model with modified signature .
    signature_model = TFModelWrapperWithSignature(
        model=model,
        preprocess_input=preprocess_input,
        input_dtype=input_dtype,
        feature_description=feature_description,
        **kwargs
        )
    tf.saved_model.save(signature_model, path)

    Step 2:
    # Load the saved_model in the beam pipeline to create ModelHandler.
    saved_model_spec = model_spec_pb2.SavedModelSpec(
        model_path=known_args.model_path)
    inferece_spec_type = model_spec_pb2.InferenceSpecType(
        saved_model_spec=saved_model_spec)
    model_handler = CreateModelHandler(inferece_spec_type)


def build_inputs_with_special_tokens(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None) -> List[int]:
    if token_ids_1 is None:
        raise ValueError('With TAPAS, you must provide both question IDs and table IDs.')
    return [self.cls_token_id] + token_ids_0 + [self.sep_token_id] + token_ids_1

Build model inputs from a question and flattened table for question answering or sequence classification tasks
    by concatenating and adding special tokens.

Args:
        token_ids_0 (`List[int]`): The ids of the question.
        token_ids_1 (`List[int]`, *optional*): The ids of the flattened table.

Returns:
        `List[int]`: The model input with special tokens.


def _SparseReorderGrad(op: ops.Operation, unused_output_indices_grad, output_values_grad):
    input_indices = op.inputs[0]
    input_shape = op.inputs[2]
    num_entries = array_ops.shape(input_indices)[0]
    entry_indices = math_ops.range(num_entries)
    sp_unordered = sparse_tensor.SparseTensor(input_indices, entry_indices, input_shape)
    sp_ordered = sparse_ops.sparse_reorder(sp_unordered)
    inverted_permutation = array_ops.invert_permutation(sp_ordered.values)
    return (None, array_ops.gather(output_values_grad, inverted_permutation), None)

Gradients for the SparseReorder op.

Args:
        op: the SparseReorder op
        unused_output_indices_grad: the incoming gradients of the output indices
        output_values_grad: the incoming gradients of the output values

Returns:
        Gradient for each of the 3 input tensors:
        (input_indices, input_values, input_shape)
        The gradients for input_indices and input_shape is None.


def write_model(model_object, output_tflite_file):
    if sys.byteorder == 'big':
        model_object = copy.deepcopy(model_object)
        byte_swap_tflite_model_obj(model_object, 'big', 'little')
    model_bytearray = convert_object_to_bytearray(model_object)
    with gfile.GFile(output_tflite_file, 'wb') as output_file_handle:
        output_file_handle.write(model_bytearray)

Writes the tflite model, a python object, into the output file.

    NOTE: This API only works for TFLite generated with
    _experimental_use_buffer_offset=false

Args:
        model_object: A tflite model as a python object
        output_tflite_file: Full path name to the output tflite file.

Raises:
        IOError: If output_tflite_file path is invalid or cannot be opened.


def lu_matrix_inverse(lower_upper, perm, validate_args=False, name=None):
    with ops.name_scope(name or 'lu_matrix_inverse'):
        lower_upper = ops.convert_to_tensor(lower_upper, dtype_hint=dtypes.float32, name='lower_upper')
        perm = ops.convert_to_tensor(perm, dtype_hint=dtypes.int32, name='perm')
        assertions = lu_reconstruct_assertions(lower_upper, perm, validate_args)
        if assertions:
            with ops.control_dependencies(assertions):
                lower_upper = array_ops.identity(lower_upper)
                perm = array_ops.identity(perm)
        shape = array_ops.shape(lower_upper)
        return lu_solve(lower_upper, perm, rhs=eye(shape[-1], batch_shape=shape[:-2], dtype=lower_upper.dtype), validate_args=False)

Computes the inverse given the LU decomposition(s) of one or more matrices.

    This op is conceptually identical to,

    ```python
    inv_X = tf.lu_matrix_inverse(*tf.linalg.lu(X))
    tf.assert_near(tf.matrix_inverse(X), inv_X)
    # ==> True
    ```

    Note: this function does not verify the implied matrix is actually invertible
    nor is this condition checked even when `validate_args=True`.

Args:
        lower_upper: `lu` as returned by `tf.linalg.lu`, i.e., if `matmul(P,
        matmul(L, U)) = X` then `lower_upper = L + U - eye`.
        perm: `p` as returned by `tf.linag.lu`, i.e., if `matmul(P, matmul(L, U)) =
        X` then `perm = argmax(P)`.
        validate_args: Python `bool` indicating whether arguments should be checked
        for correctness. Note: this function does not verify the implied matrix is
        actually invertible, even when `validate_args=True`.
        Default value: `False` (i.e., don't validate arguments).
        name: Python `str` name given to ops managed by this object.
        Default value: `None` (i.e., 'lu_matrix_inverse').

Returns:
        inv_x: The matrix_inv, i.e.,
        `tf.matrix_inverse(tf.linalg.lu_reconstruct(lu, perm))`.

    #### Examples

    ```python
    import numpy as np
    import tensorflow as tf
    import tensorflow_probability as tfp

    x = [[[3., 4], [1, 2]],
        [[7., 8], [3, 4]]]
    inv_x = tf.linalg.lu_matrix_inverse(*tf.linalg.lu(x))
    tf.assert_near(tf.matrix_inverse(x), inv_x)
    # ==> True
    ```


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

Prepare the images structure for processing.

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

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


def update_hash_with_array(hash_value, int_array):
    if int_array is not None:
        for i in int_array:
            hash_value = update_hash_with_primitive_value(hash_value, i)
    return hash_value

Update the hash value using a TFLite int array.

Args:
        hash_value (int): The current hash value.
        int_array: A TFLite int array to incorporate into the hash.

Returns:
        int: The updated hash value.


def compute_local_attention_representations(self, global_k, global_v, local_q, local_k, local_v, mask, dim: DimensionInfo):
    blocked_local_q = local_q.view(dim.batch_size, dim.num_heads, dim.num_blocks, dim.block_size, dim.dim_per_head)
    blocked_local_k = local_k.view(dim.batch_size, dim.num_heads, dim.num_blocks, dim.block_size, dim.dim_per_head)
    blocked_local_v = local_v.view(dim.batch_size, dim.num_heads, dim.num_blocks, dim.block_size, dim.dim_per_head)
    extended_mask = nn.functional.pad(mask.view(dim.batch_size, dim.num_blocks, dim.block_size), pad=(dim.global_len, 0), value=0)
    blocked_local2global = torch.einsum('BHNKF,BHGF->BHNKG', blocked_local_q, global_k)
    blocked_local2local = torch.einsum('BHNKF,BHNXF->BHNKX', blocked_local_q, blocked_local_k)
    attn_weights = torch.cat([blocked_local2global, blocked_local2local], dim=-1)
    attn_weights = attn_weights + extended_mask[:, None, :, None, :]
    attn_probs = nn.functional.softmax(attn_weights, dim=-1)
    attn_probs = nn.functional.dropout(attn_probs, p=self.dropout, training=self.training)
    local2global_attn_probs = attn_probs[:, :, :, :, :dim.global_len]
    local2local_attn_probs = attn_probs[:, :, :, :, dim.global_len:]
    local2global_attn_output = torch.einsum('BHNKG,BHGF->BHNKF', local2global_attn_probs, global_v)
    local2local_attn_output = torch.einsum('BHNKX,BHNXF->BHNKF', local2local_attn_probs, blocked_local_v)
    attn_output = local2global_attn_output + local2local_attn_output
    return (attn_output, attn_probs)

Compute attention representations for local tokens.

    Local tokens will attend to both global tokens as well as all other tokens within the same local block. Hence,
    we need to tile and concatenate the global tokens to every local block

Args:
        global_k (`torch.FloatTensor`) of shape [batch_size, num_heads, global_len, dim_per_head]:
            key vectors from global tokens
        global_v (`torch.FloatTensor`) of shape [batch_size, num_heads, global_len, dim_per_head]:
            value vectors from global tokens
        local_q (`torch.FloatTensor`) of shape [batch_size, num_heads, padded_seq_len, dim_per_head]:
            query vectors from local tokens
        local_k (`torch.FloatTensor`) of shape [batch_size, num_heads, padded_seq_len, dim_per_head]:
            key vectors from local tokens
        local_v (`torch.FloatTensor`) of shape [batch_size, num_heads, padded_seq_len, dim_per_head]:
            value vectors from local tokens
        mask (`torch.FloatTensor`) of shape [batch_size, padded_seq_len]: attention mask
        dim (DimensionInfo): DimensionInfo wrapper for dimensions

Returns:
        output of shape `[batch_sizes, length, features]`. where length will be padded to a multiple of block_size


def __getitem__(self, index):
    raise NotImplementedError

Gets batch at position `index`.

Args:
        index: position of the batch in the Sequence.

Returns:
        A batch


def lu_factor(x):
    if any_symbolic_tensors((x,)):
        return LuFactor().symbolic_call(x)
    return _lu_factor(x)

Computes the lower-upper decomposition of a square matrix.

Args:
        x: A tensor of shape `(..., M, M)`.

Returns:
        A tuple of two tensors: a tensor of shape `(..., M, M)` containing the
        lower and upper triangular matrices and a tensor of shape `(..., M)`
        containing the pivots.


def test_sample_paths_1d(self, use_batch, watch_params, supply_normal_draws, random_type):
    dtype = tf.float64
    mu = 0.2
    a = 0.4
    b = 0.33

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

    def vol_fn(t, x):
        del x
        if not use_batch:
            return (a * t + b) * tf.ones([1, 1], dtype=t.dtype)
        else:
            return (a * t + b) * tf.ones([2, 1, 1, 1], dtype=t.dtype)
    times = np.array([0.0, 0.1, 0.21, 0.32, 0.43, 0.55])
    num_samples = 10000
    if supply_normal_draws:
        normal_draws = tf.random.stateless_normal(shape=[2, 5000, 55, 1], seed=[1, 42], dtype=dtype)
        normal_draws = tf.concat([normal_draws, -normal_draws], axis=1)
    else:
        normal_draws = None
    if use_batch:
        x0 = np.array([[[0.1]], [[0.1]]])
    else:
        x0 = np.array([0.1])
    paths = self.evaluate(euler_sampling.sample(dim=1, drift_fn=drift_fn, volatility_fn=vol_fn, times=times, num_samples=num_samples, initial_state=x0, random_type=random_type, normal_draws=normal_draws, watch_params=watch_params, time_step=0.01, seed=[1, 42], dtype=dtype))
    paths_no_zero = self.evaluate(euler_sampling.sample(dim=1, drift_fn=drift_fn, volatility_fn=vol_fn, times=times[1:], num_samples=num_samples, initial_state=x0, random_type=random_type, normal_draws=normal_draws, time_step=0.01, seed=[1, 42], dtype=dtype))
    with self.subTest('CorrectShape'):
        if not use_batch:
            self.assertAllClose(paths.shape, (num_samples, 6, 1), atol=0)
        else:
            self.assertAllClose(paths.shape, (2, num_samples, 6, 1), atol=0)
    if not use_batch:
        means = np.mean(paths, axis=0).reshape(-1)
    else:
        means = np.mean(paths, axis=1).reshape([2, 1, 6])
    expected_means = x0 + 2.0 / 3.0 * mu * np.power(times, 1.5)
    with self.subTest('ExpectedResult'):
        self.assertAllClose(means, expected_means, rtol=0.01, atol=0.01)
    if not use_batch:
        with self.subTest('IncludeInitialState'):
            self.assertAllClose(paths[:, 1:, :], paths_no_zero)

Tests path properties for 1-dimentional Ito process.

    We construct the following Ito process.

    ````
    dX = mu * sqrt(t) * dt + (a * t + b) dW
    ````

    For this process expected value at time t is x_0 + 2/3 * mu * t^1.5 .

Args:
        use_batch: Test parameter to specify if we are testing the batch of Euler
        sampling.
        watch_params: Triggers custom for loop.
        supply_normal_draws: Supply normal draws.
        random_type: `RandomType` of the sampled normal draws.


def save(self, file_prefix, checkpoint_number=None, session=None, options=None):
    options = options or checkpoint_options.CheckpointOptions()
    feed_dict = {}
    use_session = not context.executing_eagerly() and (not ops.inside_function())
    if checkpoint_number:
        file_prefix = '%s-%d' % (file_prefix, checkpoint_number)
    if use_session:
        if self._object_graph_feed_tensor is None:
            with ops.device('/cpu:0'):
                self._object_graph_feed_tensor = constant_op.constant('', dtype=dtypes.string)
                self._file_prefix_feed_tensor = constant_op.constant('', dtype=dtypes.string)
        object_graph_tensor = self._object_graph_feed_tensor
        file_prefix_tensor = self._file_prefix_feed_tensor
        feed_dict[file_prefix_tensor] = file_prefix
    else:
        with ops.device('/cpu:0'):
            file_prefix_tensor = ops.convert_to_tensor(file_prefix, dtype=dtypes.string)
        object_graph_tensor = None
    if not tensor_util.is_tensor(file_prefix):
        file_io.recursive_create_dir(os.path.dirname(file_prefix))
    save_path, new_feed_additions = self._save_cached_when_graph_building(file_prefix_tensor, object_graph_tensor, options)
    if new_feed_additions:
        feed_dict.update(new_feed_additions)
    if not use_session:
        session = None
    elif session is None:
        session = get_session()
    if session:
        return session.run(save_path, feed_dict=feed_dict)
    elif use_session:
        raise RuntimeError(f'Unable to save checkpoint to "{file_prefix}" in graph mode without a default session. Please use `with tf.Session():` to create a session.')
    else:
        return save_path

Save a training checkpoint.

    The saved checkpoint includes variables created by this object and any
    Trackable objects it depends on at the time `Saver.save()` is called.

Args:
        file_prefix: A prefix to use for the checkpoint filenames
        (/path/to/directory/and_a_prefix). Names are generated based on this
        prefix and `checkpoint_number`, if provided.
        checkpoint_number: An integer variable or Tensor, used to number
        checkpoints. Typically this value is saved along with other variables in
        training checkpoints, which will happen automatically if it was created
        by `root_trackable` or one of its dependencies (via
        `Trackable._add_variable`).
        session: The session to evaluate variables in. Ignored when executing
        eagerly. If not provided when graph building, the default session is
        used.
        options: Optional `tf.train.CheckpointOptions` object.

Returns:
        The full path to the checkpoint.

Raises:
        RuntimeError: if called in V1 Graph mode without a default session.


def Py3GetFullArgSpec(fn):
    try:
        sig = inspect._signature_from_callable(fn, skip_bound_arg=True, follow_wrapper_chains=True, sigcls=inspect.Signature)
    except Exception:
        raise TypeError('Unsupported callable.')
    args = []
    varargs = None
    varkw = None
    kwonlyargs = []
    defaults = ()
    annotations = {}
    defaults = ()
    kwdefaults = {}
    if sig.return_annotation is not sig.empty:
        annotations['return'] = sig.return_annotation
    for param in sig.parameters.values():
        kind = param.kind
        name = param.name
        if kind is inspect._POSITIONAL_ONLY:
            args.append(name)
        elif kind is inspect._POSITIONAL_OR_KEYWORD:
            args.append(name)
            if param.default is not param.empty:
                defaults += (param.default,)
        elif kind is inspect._VAR_POSITIONAL:
            varargs = name
        elif kind is inspect._KEYWORD_ONLY:
            kwonlyargs.append(name)
            if param.default is not param.empty:
                kwdefaults[name] = param.default
        elif kind is inspect._VAR_KEYWORD:
            varkw = name
        if param.annotation is not param.empty:
            annotations[name] = param.annotation
    if not kwdefaults:
        kwdefaults = None
    if not defaults:
        defaults = None
    return inspect.FullArgSpec(args, varargs, varkw, defaults, kwonlyargs, kwdefaults, annotations)

A alternative to the builtin getfullargspec.

    The builtin inspect.getfullargspec uses:
    `skip_bound_args=False, follow_wrapped_chains=False`
    in order to be backwards compatible.

    This function instead skips bound args (self) and follows wrapped chains.

Args:
        fn: The function or class of interest.

Returns:
        An inspect.FullArgSpec namedtuple with the full arg spec of the function.


def load(archive_file: fhir_package.PackageSource) -> fhir_package.FhirPackage[structure_definition_pb2.StructureDefinition, search_parameter_pb2.SearchParameter, code_system_pb2.CodeSystem, value_set_pb2.ValueSet]:
    return fhir_package.FhirPackage.load(archive_file, _PRIMITIVE_HANDLER, structure_definition_pb2.StructureDefinition, search_parameter_pb2.SearchParameter, code_system_pb2.CodeSystem, value_set_pb2.ValueSet)

Instantiates and returns a new `FhirPackage` for FHIR R4.

Args:
        archive_file: The zip or tar file path or a function returning a file-like
        containing resources represented by this collection.

Returns:
        An instance of `FhirPackage`.

Raises:
        ValueError: In the event that the file or contents are invalid.


def prefer_static_broadcast_shape(shape1, shape2, name='prefer_static_broadcast_shape'):
    with ops.name_scope(name, values=[shape1, shape2]):

        def make_shape_tensor(x):
            return ops.convert_to_tensor(x, name='shape', dtype=dtypes.int32)

        def get_tensor_shape(s):
            if isinstance(s, tensor_shape.TensorShape):
                return s
            s_ = tensor_util.constant_value(make_shape_tensor(s))
            if s_ is not None:
                return tensor_shape.TensorShape(s_)
            return None

        def get_shape_tensor(s):
            if not isinstance(s, tensor_shape.TensorShape):
                return make_shape_tensor(s)
            if s.is_fully_defined():
                return make_shape_tensor(s.as_list())
            raise ValueError('Cannot broadcast from partially defined `TensorShape`.')
        shape1_ = get_tensor_shape(shape1)
        shape2_ = get_tensor_shape(shape2)
        if shape1_ is not None and shape2_ is not None:
            return array_ops.broadcast_static_shape(shape1_, shape2_)
        shape1_ = get_shape_tensor(shape1)
        shape2_ = get_shape_tensor(shape2)
        return array_ops.broadcast_dynamic_shape(shape1_, shape2_)

Convenience function which statically broadcasts shape when possible.

Args:
        shape1:  `1-D` integer `Tensor`.  Already converted to tensor!
        shape2:  `1-D` integer `Tensor`.  Already converted to tensor!
        name:  A string name to prepend to created ops.

Returns:
        The broadcast shape, either as `TensorShape` (if broadcast can be done
        statically), or as a `Tensor`.


def on_train_batch_begin(self, batch, logs=None):
    self.on_batch_begin(batch, logs=logs)

Called at the beginning of a training batch in `fit` methods.

    Subclasses should override for any actions to run.

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

Args:
        batch: Integer, index of batch within the current epoch.
        logs: Dict. Currently no data is passed to this argument for this
        method but that may change in the future.


def _cast_to_known_type(name):
    if name is None:
        return None
    return name.rstrip('.')

Canonicalizes a string representing a type if possible.

    # TODO(dbieber): Support additional canonicalization, such as string/str, and
    # boolean/bool.

Example:
        _cast_to_known_type("str.") == "str"

Args:
        name: A string representing a type, or None.

Returns:
        A canonicalized version of the type string.


def get_type_parameters(self, annot, seen=None):
    seen = seen or set()
    if annot in seen or not annot.formal:
        return []
    if isinstance(annot, mixin.NestedAnnotation):
        seen = seen | {annot}
    if isinstance(annot, abstract.TypeParameter):
        return [annot]
    elif isinstance(annot, abstract.TupleClass):
        annots = []
        for idx in range(annot.tuple_length):
            annots.extend(self.get_type_parameters(annot.formal_type_parameters[idx], seen))
        return annots
    elif isinstance(annot, mixin.NestedAnnotation):
        return sum((self.get_type_parameters(t, seen) for _, t in annot.get_inner_types()), [])
    return []

Returns all the TypeParameter instances that appear in the annotation.

    Note that if you just need to know whether or not the annotation contains
    type parameters, you can check its `.formal` attribute.

Args:
        annot: An annotation.
        seen: A seen set.


def __init__(self, max_shard_bytes, max_shards=None, bytes_per_string=16):
    """Creates a new `MaxSizePartitioner`.

Args:
  max_shard_bytes: The maximum size any given shard is allowed to be.
  max_shards: The maximum number of shards in `int` created taking
    precedence over `max_shard_bytes`.
  bytes_per_string: If the partition value is of type string, this provides
    an estimate of how large each string is.
"""
    if max_shard_bytes < 1:
        raise ValueError(f'Argument `max_shard_bytes` must be positive. Received {max_shard_bytes}')
    if max_shards and max_shards < 1:
        raise ValueError(f'Argument `max_shards` must be positive. Received {max_shards}')
    if bytes_per_string < 1:
        raise ValueError(f'Argument `bytes_per_string` must be positive. Received: {bytes_per_string}')
    self._max_shard_bytes = max_shard_bytes
    self._max_shards = max_shards
    self._bytes_per_string = bytes_per_string

Partitioner that keeps shards below `max_shard_bytes`.

    This partitioner ensures each shard has at most `max_shard_bytes`, and tries
    to allocate as few shards as possible, i.e., keeping shard size as large
    as possible.

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

Example:
    >>> partitioner = MaxSizePartitioner(max_shard_bytes=4)
    >>> partitions = partitioner(tf.TensorShape([6, 1]), tf.float32)
    >>> [6, 1]
    >>> partitioner = MaxSizePartitioner(max_shard_bytes=4, max_shards=2)
    >>> partitions = partitioner(tf.TensorShape([6, 1]), tf.float32)
    >>> [2, 1]
    >>> partitioner = MaxSizePartitioner(max_shard_bytes=1024)
    >>> partitions = partitioner(tf.TensorShape([6, 1]), tf.float32)
    >>> [1, 1]
    >>>
    >>> # use in ParameterServerStrategy
    >>> # strategy = tf.distribute.experimental.ParameterServerStrategy(
    >>> #   cluster_resolver=cluster_resolver, variable_partitioner=partitioner)


def _get_pmap_impl(f, devices, has_tpu):
    if has_tpu:
        output_is_list = [False]

        def recorder(args, kwargs, res):
            del args, kwargs
            output_is_list[0] = isinstance(res, list)
            return res
        f = _record_result_type(recorder, f)

    def tf_f(*tf_args):
        """A wrapper for `f` that takes/returns tensors."""
        np_args = _tf_to_np(tf_args)
        np_out = f(*np_args)
        return np_out
    if has_tpu:

        @polymorphic_function.function(autograph=False)
        def fn(inputs):
            res = tpu.replicate(tf_f, inputs)
            if res and isinstance(res[0], list) and (len(res[0]) == 1) and (not output_is_list[0]):
                res = [x[0] for x in res]
            return res
        return fn
    else:
        jit_tf_f = polymorphic_function.function(tf_f, autograph=False)

        @polymorphic_function.function(autograph=False)
        def fn(all_per_device_args):
            """Multi-device function with calls placed on the correct device."""
            results = []
            for per_device_args, device in zip(all_per_device_args, devices):
                with ops.device(device):
                    results.append(jit_tf_f(*per_device_args))
            return results
        return fn

This is a helper function to return the pmap impl.

Args:
        f: a function that takes ndarrays and returns ndarrays.
        devices: a list of strings; the device list.
        has_tpu: boolean; whether `devices` contains TPU devices.

Returns:
        A function that takes tensors and returns tensors.


def logdet(x):
    if any_symbolic_tensors((x,)):
        return Logdet().symbolic_call(x)
    return backend.math.logdet(x)

Computes log of the determinant of a hermitian positive definite matrix.

Args:
        x: Input matrix. It must 2D and square.

Returns:
        The natural log of the determinant of matrix.


def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor]=None):
    pixel_values = pixel_values.type(self.visual.dtype)
    image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
    return image_embeds

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.
        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
            The temporal, height and width of feature shape of each image in LLM.


def getqualifiedname(namespace, object_, max_depth=5, visited=None):
    if visited is None:
        visited = set()
    namespace = dict(namespace)
    for name in namespace:
        if object_ is namespace[name]:
            return name
    parent = tf_inspect.getmodule(object_)
    if parent is not None and parent is not object_ and (parent is not namespace):
        parent_name = getqualifiedname(namespace, parent, max_depth=0, visited=visited)
        if parent_name is not None:
            name_in_parent = getqualifiedname(parent.__dict__, object_, max_depth=0, visited=visited)
            assert name_in_parent is not None, 'An object should always be found in its owner module'
            return '{}.{}'.format(parent_name, name_in_parent)
    if max_depth:
        for name in namespace.keys():
            value = namespace[name]
            if tf_inspect.ismodule(value) and id(value) not in visited:
                visited.add(id(value))
                name_in_module = getqualifiedname(value.__dict__, object_, max_depth - 1, visited)
                if name_in_module is not None:
                    return '{}.{}'.format(name, name_in_module)
    return None

Returns the name by which a value can be referred to in a given namespace.

    If the object defines a parent module, the function attempts to use it to
    locate the object.

    This function will recurse inside modules, but it will not search objects for
    attributes. The recursion depth is controlled by max_depth.

Args:
        namespace: Dict[str, Any], the namespace to search into.
        object_: Any, the value to search.
        max_depth: Optional[int], a limit to the recursion depth when searching
        inside modules.
        visited: Optional[Set[int]], ID of modules to avoid visiting.
    Returns: Union[str, None], the fully-qualified name that resolves to the value
        o, or None if it couldn't be found.


def _normalize_feature_columns(feature_columns):
    if isinstance(feature_columns, fc_types.FeatureColumn):
        feature_columns = [feature_columns]
    if isinstance(feature_columns, collections_abc.Iterator):
        feature_columns = list(feature_columns)
    if isinstance(feature_columns, dict):
        raise ValueError('Expected feature_columns to be iterable, found dict.')
    for column in feature_columns:
        if not isinstance(column, fc_types.FeatureColumn):
            raise ValueError('Items of feature_columns must be a FeatureColumn. Given (type {}): {}.'.format(type(column), column))
    if not feature_columns:
        raise ValueError('feature_columns must not be empty.')
    name_to_column = {}
    for column in feature_columns:
        if column.name in name_to_column:
            raise ValueError('Duplicate feature column name found for columns: {} and {}. This usually means that these columns refer to same base feature. Either one must be discarded or a duplicated but renamed item must be inserted in features dict.'.format(column, name_to_column[column.name]))
        name_to_column[column.name] = column
    return sorted(feature_columns, key=lambda x: x.name)

Normalizes the `feature_columns` input.

    This method converts the `feature_columns` to list type as best as it can. In
    addition, verifies the type and other parts of feature_columns, required by
    downstream library.

Args:
        feature_columns: The raw feature columns, usually passed by users.

Returns:
        The normalized feature column list.

Raises:
        ValueError: for any invalid inputs, such as empty, duplicated names, etc.


def returns_collection(self) -> bool:
    return self.cardinality == Cardinality.COLLECTION or self.cardinality == Cardinality.CHILD_OF_COLLECTION

Indicates if the data type will evaluate to a collection.

Returns:
        True in the following circumstances
        - The data type represents an element with cardinality greater than one.
        - The data type represents an element with a cardinality less than or
        equal to one, but that element is a child of a collection and will
        evaluate to a collection. For example, the path Patient.name.use will
        return a collection, despite 'use' being a scalar, because it is a
        child of the collection, 'name.'
        False if the data type represents a scalar element whose parents are all
        also scalars.


def embed_check_categorical_event_shape(categorical_param, name='embed_check_categorical_event_shape'):
    with ops.name_scope(name, values=[categorical_param]):
        x = ops.convert_to_tensor(categorical_param, name='categorical_param')
        x_dtype = x.dtype.base_dtype
        max_event_size = _largest_integer_by_dtype(x_dtype) if x_dtype.is_floating else 0
        if max_event_size == 0:
            raise TypeError('Unable to validate size of unrecognized dtype ({}).'.format(x_dtype.name))
        try:
            x_shape_static = x.get_shape().with_rank_at_least(1)
        except ValueError:
            raise ValueError('A categorical-distribution parameter must have at least 1 dimension.')
        if tensor_shape.dimension_value(x_shape_static[-1]) is not None:
            event_size = x_shape_static.dims[-1].value
            if event_size < 2:
                raise ValueError('A categorical-distribution parameter must have at least 2 events.')
            if event_size > max_event_size:
                raise ValueError('Number of classes exceeds `dtype` precision, i.e., {} implies shape ({}) cannot exceed {}.'.format(x_dtype.name, event_size, max_event_size))
            return x
        else:
            event_size = array_ops.shape(x, name='x_shape')[-1]
            return control_flow_ops.with_dependencies([check_ops.assert_rank_at_least(x, 1, message='A categorical-distribution parameter must have at least 1 dimension.'), check_ops.assert_greater_equal(array_ops.shape(x)[-1], 2, message='A categorical-distribution parameter must have at least 2 events.'), check_ops.assert_less_equal(event_size, max_event_size, message='Number of classes exceeds `dtype` precision, i.e., {} dtype cannot exceed {} shape.'.format(x_dtype.name, max_event_size))], x)

Embeds checks that categorical distributions don't have too many classes.

    A categorical-type distribution is one which, e.g., returns the class label
    rather than a one-hot encoding.  E.g., `Categorical(probs)`.

    Since distributions output samples in the same dtype as the parameters, we
    must ensure that casting doesn't lose precision. That is, the
    `parameter.dtype` implies a maximum number of classes. However, since shape is
    `int32` and categorical variables are presumed to be indexes into a `Tensor`,
    we must also ensure that the number of classes is no larger than the largest
    possible `int32` index, i.e., `2**31-1`.

    In other words the number of classes, `K`, must satisfy the following
    condition:

    ```python
    K <= min(
        int(2**31 - 1),  # Largest float as an index.
        {
            dtypes.float16: int(2**11),   # Largest int as a float16.
            dtypes.float32: int(2**24),
            dtypes.float64: int(2**53),
        }.get(categorical_param.dtype.base_dtype, 0))
    ```

Args:
        categorical_param: Floating-point `Tensor` representing parameters of
        distribution over categories. The rightmost shape is presumed to be the
        number of categories.
        name: A name for this operation (optional).

Returns:
        categorical_param: Input `Tensor` with appropriate assertions embedded.

Raises:
        TypeError: if `categorical_param` has an unknown `dtype`.
        ValueError: if we can statically identify `categorical_param` as being too
        large (for being closed under int32/float casting).


def __init__(self, cell, residual_fn=None, **kwargs):
    super(ResidualWrapperBase, self).__init__(cell, **kwargs)
    self._residual_fn = residual_fn

Constructs a `ResidualWrapper` for `cell`.

Args:
        cell: An instance of `RNNCell`.
        residual_fn: (Optional) The function to map raw cell inputs and raw cell
        outputs to the actual cell outputs of the residual network.
        Defaults to calling nest.map_structure on (lambda i, o: i + o), inputs
        and outputs.
        **kwargs: dict of keyword arguments for base layer.


def forward(self, pixel_values: torch.FloatTensor, bool_masked_pos: torch.BoolTensor, head_mask: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[tuple, VideoMAEForPreTrainingOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    outputs = self.videomae(pixel_values, bool_masked_pos=bool_masked_pos, head_mask=head_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    sequence_output = outputs[0]
    sequence_output = self.encoder_to_decoder(sequence_output)
    batch_size, seq_len, num_channels = sequence_output.shape
    if bool_masked_pos is None:
        raise ValueError('One must provided a boolean mask ')
    expanded_position_embeddings = self.position_embeddings.expand(batch_size, -1, -1).type_as(pixel_values)
    expanded_position_embeddings = expanded_position_embeddings.detach().to(device=pixel_values.device, copy=True)
    pos_emb_visible = expanded_position_embeddings[~bool_masked_pos].reshape(batch_size, -1, num_channels)
    pos_emb_mask = expanded_position_embeddings[bool_masked_pos].reshape(batch_size, -1, num_channels)
    x_full = torch.cat([sequence_output + pos_emb_visible, self.mask_token + pos_emb_mask], dim=1)
    decoder_outputs = self.decoder(x_full, pos_emb_mask.shape[1])
    logits = decoder_outputs.logits
    loss = None
    with torch.no_grad():
        if self.config.num_channels != 3:
            frames = pixel_values
        else:
            device = pixel_values.device
            dtype = pixel_values.dtype
            mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device=device, dtype=dtype)[None, None, :, None, None]
            std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device=device, dtype=dtype)[None, None, :, None, None]
            frames = pixel_values * std + mean
        batch_size, time, num_channels, height, width = frames.shape
        tubelet_size, patch_size = (self.config.tubelet_size, self.config.patch_size)
        if self.config.norm_pix_loss:
            frames = frames.view(batch_size, time // tubelet_size, tubelet_size, num_channels, height // patch_size, patch_size, width // patch_size, patch_size)
            frames = frames.permute(0, 1, 4, 6, 2, 5, 7, 3).contiguous()
            frames = frames.view(batch_size, time // tubelet_size * height // patch_size * width // patch_size, tubelet_size * patch_size * patch_size, num_channels)
            frames_norm = (frames - frames.mean(dim=-2, keepdim=True)) / (frames.var(dim=-2, unbiased=True, keepdim=True).sqrt() + 1e-06)
            videos_patch = frames_norm.view(batch_size, time // tubelet_size * height // patch_size * width // patch_size, tubelet_size * patch_size * patch_size * num_channels)
        else:
            if self.config.num_channels != 3:
                raise ValueError("Can't unnormalize non-RGB images. Consider setting config.norm_pix_loss to False.")
            frames = frames.view(batch_size, time // tubelet_size, tubelet_size, num_channels, height // patch_size, patch_size, width // patch_size, patch_size)
            frames = frames.permute(0, 1, 4, 6, 2, 5, 7, 3).contiguous()
            videos_patch = frames.view(batch_size, time // tubelet_size * height // patch_size * width // patch_size, tubelet_size * patch_size * patch_size * num_channels)
        batch_size, _, num_channels = videos_patch.shape
        labels = videos_patch[bool_masked_pos].reshape(batch_size, -1, num_channels)
    loss_fct = MSELoss()
    loss = loss_fct(logits, labels)
    if not return_dict:
        output = (logits,) + outputs[1:]
        return (loss,) + output if loss is not None else output
    return VideoMAEForPreTrainingOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)

bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, sequence_length)`):
    Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Each video in the
    batch must have the same number of masked patches. Sequence length is `(num_frames // tubelet_size) *
    (image_size // patch_size) ** 2`.

Example:
    ```python
    >>> from transformers import AutoImageProcessor, VideoMAEForPreTraining
    >>> import numpy as np
    >>> import torch

    >>> num_frames = 16
    >>> video = list(np.random.randint(0, 256, (num_frames, 3, 224, 224)))

    >>> image_processor = AutoImageProcessor.from_pretrained("MCG-NJU/videomae-base")
    >>> model = VideoMAEForPreTraining.from_pretrained("MCG-NJU/videomae-base")

    >>> pixel_values = image_processor(video, return_tensors="pt").pixel_values

    >>> num_patches_per_frame = (model.config.image_size // model.config.patch_size) ** 2
    >>> seq_length = (num_frames // model.config.tubelet_size) * num_patches_per_frame
    >>> bool_masked_pos = torch.randint(0, 2, (1, seq_length)).bool()

    >>> outputs = model(pixel_values, bool_masked_pos=bool_masked_pos)
    >>> loss = outputs.loss
    ```


def FormatTree(tree, style_config=None, lines=None):
    style.SetGlobalStyle(style.CreateStyleFromConfig(style_config))
    comment_splicer.SpliceComments(tree)
    continuation_splicer.SpliceContinuations(tree)
    subtype_assigner.AssignSubtypes(tree)
    identify_container.IdentifyContainers(tree)
    split_penalty.ComputeSplitPenalties(tree)
    blank_line_calculator.CalculateBlankLines(tree)
    llines = pytree_unwrapper.UnwrapPyTree(tree)
    for lline in llines:
        lline.CalculateFormattingInformation()
    lines = _LineRangesToSet(lines)
    _MarkLinesToFormat(llines, lines)
    return reformatter.Reformat(_SplitSemicolons(llines), lines)

Format a parsed lib2to3 pytree.

    This provides an alternative entry point to YAPF.

Args:
        tree: (pytree.Node) The root of the pytree to format.
        style_config: (string) Either a style name or a path to a file that contains
        formatting style settings. If None is specified, use the default style
        as set in style.DEFAULT_STYLE_FACTORY
        lines: (list of tuples of integers) A list of tuples of lines, [start, end],
        that we want to format. The lines are 1-based indexed. It can be used by
        third-party code (e.g., IDEs) when reformatting a snippet of code rather
        than a whole file.

Returns:
        The source formatted according to the given formatting style.


def get_keras_mask(x):
    return get_tensor_attr(x, '_keras_mask')

Gets the Keras mask attribute from the given tensor.

Args:
        x: Input tensor.

Returns:
        The mask tensor associated with the input tensor, or `None` if no mask
        has been set.


def build_graph(device, n, m, k, transpose_a, transpose_b, dtype):
    with ops.device('%s' % device):
        if not transpose_a:
            x = variable_v1.VariableV1(random_ops.random_uniform([n, m], dtype=dtype), use_resource=False)
        else:
            x = variable_v1.VariableV1(random_ops.random_uniform([m, n], dtype=dtype), use_resource=False)
        if not transpose_b:
            y = variable_v1.VariableV1(random_ops.random_uniform([m, k], dtype=dtype), use_resource=False)
        else:
            y = variable_v1.VariableV1(random_ops.random_uniform([k, m], dtype=dtype), use_resource=False)
        z = math_ops.matmul(x, y, transpose_a=transpose_a, transpose_b=transpose_b)
        return control_flow_ops.group(z)

Build a graph containing a sequence of matmul operations.

Args:
        device: String, the device to run on.
        n: tensor A's first dimension size.
        m: tensor A's second dimension size.
        k: tensor B's second dimension size.
        transpose_a: boolean value to show if tensor A is transposed.
        transpose_b: boolean value to show if tensor B is transposed.
        dtype: numpy data type of the input tensor.

Returns:
        A matmul operation to run()


def random_masking(self, sequence: tf.Tensor, noise: tf.Tensor | None=None):
    batch_size, seq_length, dim = shape_list(sequence)
    len_keep = int(seq_length * (1 - self.config.mask_ratio))
    if noise is None:
        noise = tf.random.uniform(shape=(batch_size, seq_length), minval=0.0, maxval=1.0)
    ids_shuffle = tf.argsort(noise, axis=1)
    ids_restore = tf.argsort(ids_shuffle, axis=1)
    ids_keep = ids_shuffle[:, :len_keep]
    sequence_unmasked = tf.gather(sequence, axis=1, batch_dims=1, indices=ids_keep)
    mask_keep = tf.zeros((batch_size, len_keep))
    mask_remove = tf.ones((batch_size, seq_length - len_keep))
    mask = tf.concat([mask_keep, mask_remove], axis=-1)
    mask = tf.gather(mask, axis=1, batch_dims=1, indices=ids_restore)
    return (sequence_unmasked, mask, ids_restore)

Perform per-sample random masking by per-sample shuffling. Per-sample shuffling is done by argsort random
    noise.

Args:
        sequence (`tf.Tensor` of shape `(batch_size, sequence_length, dim)`)
        noise (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*) which is
            mainly used for testing purposes to control randomness and maintain the reproducibility


def sheets_tab_id(config, auth, sheet_url_or_name, sheet_tab):
    sheet_id = None
    tab_id = None
    spreadsheet = sheets_get(config, auth, sheet_url_or_name)
    if spreadsheet:
        sheet_id = spreadsheet['spreadsheetId']
        for tab in spreadsheet.get('sheets', []):
            if tab['properties']['title'] == sheet_tab:
                tab_id = tab['properties']['sheetId']
                break
    return (sheet_id, tab_id)

Pull sheet tab id from URL, name, or id itself.

Args:
        config - see starthinker/util/configuration.py
        auth - user or service
        url_or_name - one of: URL, document title, or id
        sheet_tab - name of tab to get id for

Returns:
        Pair of sheet id and tab id.


def to_pil_image(self, image, rescale=None):
    self._ensure_format_supported(image)
    if is_torch_tensor(image):
        image = image.numpy()
    if isinstance(image, np.ndarray):
        if rescale is None:
            rescale = isinstance(image.flat[0], np.floating)
        if image.ndim == 3 and image.shape[0] in [1, 3]:
            image = image.transpose(1, 2, 0)
        if rescale:
            image = image * 255
        image = image.astype(np.uint8)
        return PIL.Image.fromarray(image)
    return image

Converts `image` to a PIL Image. Optionally rescales it and puts the channel dimension back as the last axis if
    needed.

Args:
        image (`PIL.Image.Image` or `numpy.ndarray` or `torch.Tensor`):
            The image to convert to the PIL Image format.
        rescale (`bool`, *optional*):
            Whether or not to apply the scaling factor (to make pixel values integers between 0 and 255). Will
            default to `True` if the image type is a floating type, `False` otherwise.


def starting_wall_time(self):
    return self._starting_wall_time

Get the starting timestamp of the instrumented TensorFlow program.

    When there are multiple hosts (i.e., multiple tfdbg file sets), the earliest
    timestamp among the file sets is returned. It is assumed to be the job that
    starts first (e.g., the coordinator).

Returns:
        Starting timestamp in seconds since the epoch, as a float.


def before_run(self, run_context):
    if not self._grpc_debug_wrapper_session:
        self._grpc_debug_wrapper_session = grpc_wrapper.GrpcDebugWrapperSession(run_context.session, self._grpc_debug_server_addresses, watch_fn=self._watch_fn, thread_name_filter=self._thread_name_filter)
    fetches = run_context.original_args.fetches
    feed_dict = run_context.original_args.feed_dict
    watch_options = self._watch_fn(fetches, feed_dict)
    run_options = config_pb2.RunOptions()
    debug_utils.watch_graph(run_options, run_context.session.graph, debug_urls=self._grpc_debug_wrapper_session.prepare_run_debug_urls(fetches, feed_dict), debug_ops=watch_options.debug_ops, node_name_regex_allowlist=watch_options.node_name_regex_allowlist, op_type_regex_allowlist=watch_options.op_type_regex_allowlist, tensor_dtype_regex_allowlist=watch_options.tensor_dtype_regex_allowlist, tolerate_debug_op_creation_failures=watch_options.tolerate_debug_op_creation_failures)
    return session_run_hook.SessionRunArgs(None, feed_dict=None, options=run_options)

Called right before a session is run.

Args:
        run_context: A session_run_hook.SessionRunContext. Encapsulates
        information on the run.

Returns:
        A session_run_hook.SessionRunArgs object.


def split(pcoll, regex, outputEmpty=False):
    regex = Regex._regex_compile(regex)
    outputEmpty = bool(outputEmpty)

    def _process(element):
        r = regex.split(element)
        if r and (not outputEmpty):
            r = list(filter(None, r))
        yield r
    return pcoll | FlatMap(_process)

Returns the list string which was splitted on the basis of regular
    expression. It will not output empty items (by defaults).

Args:
        regex: the regular expression string or (re.compile) pattern.
        outputEmpty: (optional) Should empty be output. True to output empties
        and false if not.


def functions_to_serialize(self, serialization_cache):
    raise NotImplementedError

Returns extra functions to include when serializing a Keras object.

    Normally, when calling exporting an object to SavedModel, only the
    functions and objects defined by the user are saved. For example:

    ```
    obj = tf.Module()
    obj.v = tf.Variable(1.)

    @tf.function
    def foo(...): ...

    obj.foo = foo

    w = tf.Variable(1.)

    tf.saved_model.save(obj, 'path/to/saved/model')
    loaded = tf.saved_model.load('path/to/saved/model')

    loaded.v  # Variable with the same value as obj.v
    loaded.foo  # Equivalent to obj.foo
    loaded.w  # AttributeError
    ```

    Assigning trackable objects to attributes creates a graph, which is used for
    both checkpointing and SavedModel serialization.

    When the graph generated from attribute tracking is insufficient, extra
    objects and functions may be added at serialization time. For example,
    most models do not have their call function wrapped with a @tf.function
    decorator. This results in `model.call` not being saved. Since Keras objects
    should be revivable from the SavedModel format, the call function is added
    as an extra function to serialize.

    This function and `objects_to_serialize` is called multiple times when
    exporting to SavedModel. Please use the cache to avoid generating new
    functions and objects. A fresh cache is created for each SavedModel export.

Args:
        serialization_cache: Dictionary passed to all objects in the same object
        graph during serialization.

Returns:
        A dictionary mapping attribute names to `Function` or
        `ConcreteFunction`.


def __init__(self, options, autograph_module=None):
    self.options = options
    self.autograph_module = autograph_module

ProgramContext keeps track of converting function hierarchies.

Attributes:
        options: ConversionOptions
        autograph_module: Deprecated. Do not use.


def _MaxPoolAlongRows(self, input_matrix, row_seq, overlapping):
    output_image = np.zeros(input_matrix.shape[1])
    row_max = row_seq[-1]
    for i in range(row_seq.shape[0] - 1):
        row_start = row_seq[i]
        row_end = row_seq[i + 1] + 1 if overlapping else row_seq[i + 1]
        row_end = min(row_end, row_max)
        output_image = np.vstack((output_image, np.amax(input_matrix[row_start:row_end, :], axis=0)))
    return output_image[1:, :]

Perform max pool along row of a 2-D matrix based on row_seq.

Args:
        input_matrix: A 2-D matrix.
        row_seq: Cumulative pooling sequence along row.
        overlapping: Whether or not use overlapping when pooling.

Returns:
        A 2-D matrix, with
        * num_rows = len(row_seq)-1
        * num_cols = input_matrix.num_cols.


def account_displayed_op_only(self, is_true):
    self._options['account_displayed_op_only'] = is_true
    return self

Whether only account the statistics of displayed profiler nodes.

Args:
        is_true: If true, only account statistics of nodes eventually
        displayed by the outputs.
        Otherwise, a node's statistics are accounted by its parents
        as long as it's types match 'account_type_regexes', even if
        it is hidden from the output, say, by hide_name_regexes.

Returns:
        self