def report_list(config, auth, account):
    account_id, advertiser_id = parse_account(config, auth, account)
    is_superuser, profile_id = get_profile_for_api(config, auth, account_id)
    kwargs = {'profileId': profile_id, 'accountId': account_id} if is_superuser else {'profileId': profile_id}
    for report in API_DCM(auth, iterate=True, internal=is_superuser).reports().list(**kwargs).execute():
        yield report

Lists all the DCM report configurations for an account given the current credentials.

    Bulletproofing:
    https://developers.google.com/doubleclick-advertisers/v3.2/reports/list

Args:
        * auth: (string) Either user or service.
        * account: (string) [account:advertiser@profile] token.

Returns:
        * Iterator of JSONs.


def forward(self, pixel_values: Tensor, pixel_mask: Optional[Tensor]=None, output_hidden_states: Optional[bool]=None, output_attentions: Optional[bool]=None, return_dict: Optional[bool]=None) -> MaskFormerModelOutput:
    if pixel_values is None:
        raise ValueError('You have to specify pixel_values')
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    batch_size, _, height, width = pixel_values.shape
    if pixel_mask is None:
        pixel_mask = torch.ones((batch_size, height, width), device=pixel_values.device)
    pixel_level_module_output = self.pixel_level_module(pixel_values, output_hidden_states, return_dict=return_dict)
    image_features = pixel_level_module_output[0]
    pixel_embeddings = pixel_level_module_output[1]
    transformer_module_output = self.transformer_module(image_features, output_hidden_states, output_attentions)
    queries = transformer_module_output.last_hidden_state
    encoder_hidden_states = None
    pixel_decoder_hidden_states = None
    transformer_decoder_hidden_states = None
    hidden_states = None
    if output_hidden_states:
        encoder_hidden_states = pixel_level_module_output[2]
        pixel_decoder_hidden_states = pixel_level_module_output[3]
        transformer_decoder_hidden_states = transformer_module_output[1]
        hidden_states = encoder_hidden_states + pixel_decoder_hidden_states + transformer_decoder_hidden_states
    output = MaskFormerModelOutput(encoder_last_hidden_state=image_features, pixel_decoder_last_hidden_state=pixel_embeddings, transformer_decoder_last_hidden_state=queries, encoder_hidden_states=encoder_hidden_states, pixel_decoder_hidden_states=pixel_decoder_hidden_states, transformer_decoder_hidden_states=transformer_decoder_hidden_states, hidden_states=hidden_states, attentions=transformer_module_output.attentions)
    if not return_dict:
        output = tuple((v for v in output.values()))
    return output

Examples:

    ```python
    >>> from transformers import AutoImageProcessor, MaskFormerModel
    >>> from PIL import Image
    >>> import requests

    >>> # load MaskFormer fine-tuned on ADE20k semantic segmentation
    >>> image_processor = AutoImageProcessor.from_pretrained("facebook/maskformer-swin-base-ade")
    >>> model = MaskFormerModel.from_pretrained("facebook/maskformer-swin-base-ade")

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

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

    >>> # forward pass
    >>> outputs = model(**inputs)

    >>> # the decoder of MaskFormer outputs hidden states of shape (batch_size, num_queries, hidden_size)
    >>> transformer_decoder_last_hidden_state = outputs.transformer_decoder_last_hidden_state
    >>> list(transformer_decoder_last_hidden_state.shape)
    [1, 100, 256]
    ```


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

Returns the index of the maximum value along an axis.

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

Returns:
        A tensor.


def forward_rates(df_start_dates, df_end_dates, daycount_fractions, dtype=None, name=None):
    name = name or 'forward_rates'
    with tf.name_scope(name):
        df_start_dates = tf.convert_to_tensor(df_start_dates, dtype, name='df_start_dates')
        dtype = dtype or df_start_dates.dtype
        df_end_dates = tf.convert_to_tensor(df_end_dates, dtype, name='df_end_dates')
        daycount_fractions = tf.convert_to_tensor(daycount_fractions, dtype, name='daycount_fractions')
        return tf.math.divide_no_nan(tf.math.divide_no_nan(df_start_dates, df_end_dates) - 1, daycount_fractions)

Computes forward rates from daycount fractions and discount factors.

    #### Example
    ```python
    # Discount factors at start dates
    df_start_dates = [[0.95, 0.9, 0.75], [0.95, 0.99, 0.85]]
    # Discount factors at end dates
    df_end_dates = [[0.8, 0.6, 0.5], [0.8, 0.9, 0.5]]
    # Daycount fractions between the dates
    daycount_fractions = [[0.5, 1.0, 2], [0.6, 0.4, 4.0]]
    # Expected:
    #  [[0.375 , 0.5   , 0.25  ],
    #   [0.3125, 0.25  , 0.175 ]]
    forward_rates(df_start_dates, df_end_dates, daycount_fractions,
                dtype=tf.float64)
    ```

Args:
        df_start_dates: A real `Tensor` representing discount factors at the start
        dates.
        df_end_dates: A real `Tensor` representing discount factors at the end
        dates.
        daycount_fractions: A real `Tensor` representing  year fractions for the
        coupon accrual.
        dtype: `tf.Dtype`. If supplied the dtype for the input and output `Tensor`s.
        Default value: None which maps to the default dtype inferred from
        `df_start_dates`.
        name: Python str. The name to give to the ops created by this function.
        Default value: None which maps to 'forward_rates'.

    Returns:


def executions(self, digest=False, begin=None, end=None):
    digests = self._execution_digests
    if begin is not None or end is not None:
        begin = begin or 0
        end = end or len(digests)
        digests = digests[begin:end]
    if digest:
        return digests
    else:
        return [self.read_execution(digest) for digest in digests]

Get `Execution`s or `ExecutionDigest`s this reader has read so far.

Args:
        digest: Whether the results are returned in a digest form, i.e.,
        `ExecutionDigest` format, instead of the more detailed `Execution`
        format.
        begin: Optional beginning index for the requested execution data objects
        or their digests. Python-style negative indices are supported.
        end: Optional ending index for the requested execution data objects or
        their digests. Python-style negative indices are supported.

Returns:
        If `digest`: a `list` of `ExecutionDigest` objects.
        Else: a `list` of `Execution` objects.


def np_asarray(values, dtype=None, order=None, copy=None):
    if np.lib.NumpyVersion(np.__version__) >= '2.0.0.dev0':
        if dtype is not None and np.issubdtype(dtype, np.number):
            return np.asarray(values, order=order, copy=copy).astype(dtype, copy=copy)
        else:
            return np.asarray(values, dtype=dtype, order=order, copy=copy)
    else:
        return np.asarray(values, dtype=dtype, order=order)

Converts input values to a NumPy array.

    It will not make a copy.

    In NumPy 2.x and later, strict type casting can lead to errors when values
    overflow the specified dtype. This function addresses this by replacing direct
    np.array(..., dtype=...) calls with np.array(...).astype(...). This allows for
    intended overflows, aligning with the behavior of older NumPy versions.

Args:
        values: Array_like objects. E.g., a python list, tuple, or an object whose
        __array__ method returns an array.
        dtype: The desired numpy data type for the array.
        order: {‘C’, ‘F’, ‘A’, ‘K’}.
        copy: bool. If True, then the object is copied. If None then the object is
        copied only if needed, i.e. if __array__ returns a copy, if obj is a
        nested sequence, or if a copy is needed to satisfy any of the other
        requirements (dtype, order, etc.). For False it raises a ValueError if a
        copy cannot be avoided.

Returns:
        A NumPy array with the specified data type.


def is_layouts_same(self, embedding_layouts) -> bool:
    if self._checkpoint_layouts.keys() != embedding_layouts.keys():
        raise ValueError('Layouts in checkpoint and embedding must have the same keys. found {} and {}'.format(self._checkpoint_layouts.keys(), embedding_layouts.keys()))
    for key, layout in self._checkpoint_layouts.items():
        if not compare.ProtoEq(layout, embedding_layouts[key]):
            logging.info('Layouts do not match for %s this will require resharding; %s vs %s', key, layout, embedding_layouts[key])
            return False
    return True

Returns True if the all the embedding and checkpoint layouts are the same.

Args:
        embedding_layouts: dict of layouts for embedding tables.

    Raises: ValueError if the embedding layouts and checkpoint layouts do not
        have the same keys.
    Returns: Bool representing if the embedding layouts match the layouts in
        checkpoint.


def _set_verbosity_from(posarg):

    def decorator(f):

        def wrapper(*args, **kwargs):
            options = kwargs.get('options', args[posarg])
            with config.verbosity_from(options):
                return f(*args, **kwargs)
        return wrapper
    return decorator

Decorator to set the verbosity for a function that takes an options arg.

    Assumes that the function has an argument named `options` that is a
    config.Options object.

Args:
        posarg: The index of `options` in the positional arguments.

Returns:
        The decorator.


def stddev(self, name='stddev'):
    with self._name_scope(name):
        try:
            return self._stddev()
        except NotImplementedError as original_exception:
            try:
                return math_ops.sqrt(self._variance())
            except NotImplementedError:
                raise original_exception

Standard deviation.

    Standard deviation is defined as,

    ```none
    stddev = E[(X - E[X])**2]**0.5
    ```

    where `X` is the random variable associated with this distribution, `E`
    denotes expectation, and `stddev.shape = batch_shape + event_shape`.

Args:
        name: Python `str` prepended to names of ops created by this function.

Returns:
        stddev: Floating-point `Tensor` with shape identical to
        `batch_shape + event_shape`, i.e., the same shape as `self.mean()`.


def summarize(self, highlight=None):
    lines = [RL('Command-line configuration:', 'bold'), RL('')]
    for name, val in self._config.items():
        highlight_attr = 'bold' if name == highlight else None
        line = RL('  ')
        line += RL(name, ['underline', highlight_attr])
        line += RL(': ')
        line += RL(str(val), font_attr=highlight_attr)
        lines.append(line)
    return debugger_cli_common.rich_text_lines_from_rich_line_list(lines)

Get a text summary of the config.

Args:
        highlight: A property name to highlight in the output.

Returns:
        A `RichTextLines` output.


def merge_summaries(prev_summary, next_summary, epsilon):
    merged = np.concatenate((prev_summary, next_summary), axis=1)
    merged = np.take(merged, np.argsort(merged[0]), axis=1)
    return compress_summary(merged, epsilon)

Weighted merge sort of summaries.

    Given two summaries of distinct data, this function merges (and compresses)
    them to stay within `epsilon` error tolerance.

Args:
        prev_summary: 2D `np.ndarray` summary to be merged with `next_summary`.
        next_summary: 2D `np.ndarray` summary to be merged with `prev_summary`.
        epsilon: A float that determines the approximate desired precision.

Returns:
        A 2-D `np.ndarray` that is a merged summary. First column is the
        interpolated partition values, the second is the weights (counts).


def split(self, path):
    return os.path.split(os.path.abspath(path))

Splits the given path into two parts.

    Splits the path into a pair (head, tail) such that tail contains the last
    component of the path and head contains everything up to that.

Args:
        path: path as a string

Returns:
        a pair of path components as strings.


def delete_recursively_v2(path):
    _pywrap_file_io.DeleteRecursively(compat.path_to_bytes(path))

Deletes everything under path recursively.

Args:
        path: string, a path

Raises:
        errors.OpError: If the operation fails.


def eq(left: Any, right: Any) -> bool:
    if left is right:
        return True
    if isinstance(left, list) and isinstance(right, list) or (isinstance(left, tuple) and isinstance(right, tuple)):
        if len(left) != len(right):
            return False
        for x, y in zip(left, right):
            if ne(x, y):
                return False
        return True
    elif isinstance(left, dict):
        if not isinstance(right, dict) or len(left) != len(right) or set(left.keys()) != set(right.keys()):
            return False
        left_items = left.sym_items if isinstance(left, Symbolic) else left.items
        right_item = right.sym_getattr if isinstance(right, Symbolic) else right.__getitem__
        for k, v in left_items():
            if ne(v, right_item(k)):
                return False
        return True
    elif hasattr(left, 'sym_eq') and (not inspect.isclass(left)) and (left.sym_eq.__code__ is not Symbolic.sym_eq.__code__):
        return left.sym_eq(right)
    elif hasattr(right, 'sym_eq') and (not inspect.isclass(right)) and (right.sym_eq.__code__ is not Symbolic.sym_eq.__code__):
        return right.sym_eq(left)
    return pg_typing.callable_eq(left, right)

Compares if two values are equal. Use symbolic equality if possible.

    Example::

        @pg.members([
        ('x', pg.typing.Any())
        ])
        class A(pg.Object):
        def sym_eq(self, right):
        if super().sym_eq(right):
            return True
        return pg.eq(self.x, right)

        class B:
        pass

        assert pg.eq(1, 1)
        assert pg.eq(A(1), A(1))
        # This is True since A has override `sym_eq`.
        assert pg.eq(A(1), 1)
        # Objects of B are compared by references.
        assert not pg.eq(A(B()), A(B()))

Args:
        left: The left-hand value to compare.
        right: The right-hand value to compare.

Returns:
        True if left and right is equal or symbolically equal. Otherwise False.


def RegisterImplementation(cache_name, map_name, cache):
    global _cache_implementations
    if cache_name not in _cache_implementations:
        logging.info('Registering [%s] cache for [%s].', cache_name, map_name)
        _cache_implementations[cache_name] = {}
    _cache_implementations[cache_name][map_name] = cache

Register a Cache implementation with the CacheFactory.

    Child modules are expected to call this method in the file-level scope
    so that the CacheFactory is aware of them.

Args:
        cache_name: (string) The name of the NSS backend.
        map_name: (string) The name of the map handled by this Cache.
        cache: A class type that is a subclass of Cache.

    Returns: Nothing


def get_resource(self, feature_column, name):
    del feature_column, name
    raise NotImplementedError('StateManager.get_resource')

Returns an already created resource.

    Resources can be things such as tables, variables, trackables, etc.

Args:
        feature_column: A `FeatureColumn` object this variable corresponds to.
        name: Name of the resource.


def rescale(self, image: np.ndarray, rescale_factor: float, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> np.ndarray:
    return rescale(image, rescale_factor, data_format=data_format, input_data_format=input_data_format)

Rescale the image by the given factor. image = image * rescale_factor.

Args:
        image (`np.ndarray`):
            Image to rescale.
        rescale_factor (`float`):
            The value to use for rescaling.
        data_format (`str` or `ChannelDimension`, *optional*):
            The channel dimension format for the output image. If unset, the channel dimension format of the input
            image is used. Can be one of:
            - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
            - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
        input_data_format (`str` or `ChannelDimension`, *optional*):
            The channel dimension format for the input image. If unset, is inferred from the input image. Can be
            one of:
            - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
            - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.


def invert_attention_mask(self, encoder_attention_mask: Tensor) -> Tensor:
    if encoder_attention_mask.dim() == 3:
        encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]
    if encoder_attention_mask.dim() == 2:
        encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]
    encoder_extended_attention_mask = encoder_extended_attention_mask.to(dtype=self.dtype)
    encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * torch.finfo(self.dtype).min
    return encoder_extended_attention_mask

Invert an attention mask (e.g., switches 0. and 1.).

Args:
        encoder_attention_mask (`torch.Tensor`): An attention mask.

Returns:
        `torch.Tensor`: The inverted attention mask.


def get_use_xla_spmd(device_type):
    return device_type == 'TPU' and '0' != os.environ.get('DTENSOR_TEST_USE_XLA_SPMD', '0')

Returns True when device_type is TPU and environment variable is set.

Args:
        device_type: A str representing the type of device on the mesh.

Returns:
        bool: True when device_type is TPU and environment variable is set.


def conv_output_shape(input_shape, kernel_shape, strides, padding):
    dims = range(len(kernel_shape))
    output_shape = [conv_output_length(input_shape[d], kernel_shape[d], padding, strides[d]) for d in dims]
    output_shape = tuple([0 if input_shape[d] == 0 else output_shape[d] for d in dims])
    return output_shape

Return the output shape of an N-D convolution.

    Forces dimensions where input is empty (size 0) to remain empty.

Args:
        input_shape: tuple of size N: `(d_in1, ..., d_inN)`, spatial shape of the
        input.
        kernel_shape: tuple of size N, spatial shape of the convolutional kernel /
        receptive field.
        strides: tuple of size N, strides along each spatial dimension.
        padding: type of padding, string `"same"` or `"valid"`.
        `"valid"` means no padding. `"same"` results in padding evenly to
        the left/right or up/down of the input such that output has the same
        height/width dimension as the input.

Returns:
        tuple of size N: `(d_out1, ..., d_outN)`, spatial shape of the output.


def get_layer(self, name=None, index=None):
    if index is not None and name is not None:
        raise ValueError('Provide only a layer name or a layer index.')
    if index is not None:
        if len(self.layers) <= index:
            raise ValueError('Was asked to retrieve layer at index ' + str(index) + ' but model only has ' + str(len(self.layers)) + ' layers.')
        else:
            return self.layers[index]
    if name is not None:
        for layer in self.layers:
            if layer.name == name:
                return layer
        raise ValueError('No such layer: ' + name + '.')
    raise ValueError('Provide either a layer name or layer index.')

Retrieves a layer based on either its name (unique) or index.

    If `name` and `index` are both provided, `index` will take precedence.
    Indices are based on order of horizontal graph traversal (bottom-up).

Args:
        name: String, name of layer.
        index: Integer, index of layer.

Returns:
        A layer instance.

Raises:
        ValueError: In case of invalid layer name or index.


def __init__(self, config: FastSpeech2ConformerConfig, num_layers=2, num_chans=384, kernel_size=3, dropout_rate=0.5):
    super().__init__()
    self.conv_layers = nn.ModuleList()
    for idx in range(num_layers):
        input_channels = config.hidden_size if idx == 0 else num_chans
        layer = FastSpeech2ConformerPredictorLayer(input_channels, num_chans, kernel_size, dropout_rate)
        self.conv_layers.append(layer)
    self.linear = nn.Linear(num_chans, 1)

Initialize variance predictor module.

Args:
        input_dim (`int`): Input dimension.
        num_layers (`int`, *optional*, defaults to 2): Number of convolutional layers.
        num_chans (`int`, *optional*, defaults to 384): Number of channels of convolutional layers.
        kernel_size (`int`, *optional*, defaults to 3): Kernel size of convolutional layers.
        dropout_rate (`float`, *optional*, defaults to 0.5): Dropout rate.


def _get_saver_def_or_none(exported_model: exported_model_pb2.ExportedModel) -> Optional[saver_pb2.SaverDef]:
    if exported_model.HasField('saver_def'):
        return exported_model.saver_def
    return None

Returns the SaverDef from ExportedModel, None otherwise.

Args:
        exported_model: ExportedModel to take the SaverDef from.

Returns:
        SaverDef instance if the field `saver_def` is set. None otherwise.


def get_attr(self, name):
    fields = ('s', 'i', 'f', 'b', 'type', 'shape', 'tensor', 'func')
    try:
        with c_api_util.tf_buffer() as buf:
            pywrap_tf_session.TF_OperationGetAttrValueProto(self._c_op, name, buf)
            data = pywrap_tf_session.TF_GetBuffer(buf)
    except errors.InvalidArgumentError as e:
        raise ValueError(e.message)
    x = attr_value_pb2.AttrValue()
    x.ParseFromString(data)
    oneof_value = x.WhichOneof('value')
    if oneof_value is None:
        return []
    if oneof_value == 'list':
        for f in fields:
            if getattr(x.list, f):
                if f == 'type':
                    return [dtypes.as_dtype(t) for t in x.list.type]
                else:
                    return list(getattr(x.list, f))
        return []
    if oneof_value == 'type':
        return dtypes.as_dtype(x.type)
    assert oneof_value in fields, 'Unsupported field type in ' + str(x)
    return getattr(x, oneof_value)

Returns the value of the attr of this op with the given `name`.

Args:
        name: The name of the attr to fetch.

Returns:
        The value of the attr, as a Python object.

Raises:
        ValueError: If this op does not have an attr with the given `name`.


def _extract_type_spec_recursively(value):
    if isinstance(value, composite_tensor.CompositeTensor):
        return value._type_spec
    if isinstance(value, variables.Variable):
        return resource_variable_ops.VariableSpec(value.shape, dtype=value.dtype, trainable=value.trainable)
    if tensor_util.is_tensor(value):
        return tensor_spec.TensorSpec(value.shape, value.dtype)
    if isinstance(value, list):
        return list((_extract_type_spec_recursively(v) for v in value))
    if isinstance(value, data_structures.TrackableDataStructure):
        return _extract_type_spec_recursively(value.__wrapped__)
    if isinstance(value, tuple):
        return type(value)((_extract_type_spec_recursively(x) for x in value))
    if isinstance(value, dict):
        return type(value)(((k, _extract_type_spec_recursively(v)) for k, v in value.items()))
    return value

Return (collection of) `TypeSpec`(s) for `value` if it includes `Tensor`s.

    If `value` is a `Tensor` or `CompositeTensor`, return its `TypeSpec`. If
    `value` is a collection containing `Tensor` values, recursively supplant them
    with their respective `TypeSpec`s in a collection of parallel stucture.

    If `value` is none of the above, return it unchanged.

Args:
        value: a Python `object` to (possibly) turn into a (collection of)
        `tf.TypeSpec`(s).

Returns:
        spec: the `TypeSpec` or collection of `TypeSpec`s corresponding to `value`
        or `value`, if no `Tensor`s are found.


def concatenate(xs, axis=0):
    if any_symbolic_tensors(xs):
        return Concatenate(axis=axis).symbolic_call(xs)
    return backend.numpy.concatenate(xs, axis=axis)

Join a sequence of tensors along an existing axis.

Args:
        xs: The sequence of tensors to concatenate.
        axis: The axis along which the tensors will be joined. Defaults to `0`.

Returns:
        The concatenated tensor.


def get_registered_kernels_for_op(name):
    buf = c_api.TF_GetRegisteredKernelsForOp(name)
    data = c_api.TF_GetBuffer(buf)
    kernel_list = kernel_def_pb2.KernelList()
    kernel_list.ParseFromString(compat.as_bytes(data))
    return kernel_list

Returns a KernelList proto of registered kernels for a given op.

Args:
        name: A string representing the name of the op whose kernels to retrieve.


def _determine_and_instrument_traced_tensors(self, graph_order, ops_in_exec_path, tensor_trace_points, report_handler):
    traced_tensors = []
    checkpoint_operations = set([tensor.op for tensor, _ in tensor_trace_points])
    for op_id, op in enumerate(graph_order.operations):
        if checkpoint_operations and op not in checkpoint_operations:
            continue
        if self._skip_op(op_id, op, ops_in_exec_path, report_handler):
            continue
        for i in range(len(op.outputs)):
            out_tensor = op.outputs[i]
            if not self._skip_tensor(op_id, out_tensor, report_handler):
                traced_tensors.append(out_tensor)
    return traced_tensors

Determines the tensors to trace and instruments the trace details.

Args:
        graph_order: graph_order tuple containing graph (tf.graph), operations
        (list of operations), op_to_idx (op id mapping), (tensors) list of
        tensors, tensor_to_idx (tensor id mapping), contains_cycle (whether
        there is a cycle in the graph), topological_order_or_cycle (list of ops
        in topological order or list of ops creating a cycle).
        ops_in_exec_path: Set of ops in the execution path.
        tensor_trace_points: Collection of programatic tensor trace points.
        report_handler: An instance of tensor_tracer_report.TTReportHandle.

Returns:
        List of tensors to be traced.


def _get_profile_data_generator(self):
    node_to_file_path = {}
    node_to_line_number = {}
    node_to_func_name = {}
    node_to_op_type = {}
    for op in self._graph.get_operations():
        for trace_entry in reversed(op.traceback):
            file_path = trace_entry[0]
            line_num = trace_entry[1]
            func_name = trace_entry[2]
            if not source_utils.guess_is_tensorflow_py_library(file_path):
                break
        node_to_file_path[op.name] = file_path
        node_to_line_number[op.name] = line_num
        node_to_func_name[op.name] = func_name
        node_to_op_type[op.name] = op.type

    def profile_data_generator(device_step_stats):
        for node_stats in device_step_stats.node_stats:
            if node_stats.node_name == '_SOURCE' or node_stats.node_name == '_SINK':
                continue
            yield profiling.ProfileDatum(device_step_stats.device, node_stats, node_to_file_path.get(node_stats.node_name, ''), node_to_line_number.get(node_stats.node_name, 0), node_to_func_name.get(node_stats.node_name, ''), node_to_op_type.get(node_stats.node_name, ''))
    return profile_data_generator

Get function that generates `ProfileDatum` objects.

Returns:
        A function that generates `ProfileDatum` objects.


def __init__(self, dump_root, tfdbg_run_id, circular_buffer_size=DEFAULT_CIRCULAR_BUFFER_SIZE):
    if not dump_root:
        raise ValueError('Empty or None dump root')
    self._dump_root = dump_root
    self._tfdbg_run_id = tfdbg_run_id
    _pywrap_debug_events_writer.Init(self._dump_root, self._tfdbg_run_id, circular_buffer_size)

Construct a DebugEventsWriter object.

    NOTE: Given the same `dump_root`, all objects from this constructor
        will point to the same underlying set of writers. In other words, they
        will write to the same set of debug events files in the `dump_root`
        folder.

Args:
        dump_root: The root directory for dumping debug data. If `dump_root` does
        not exist as a directory, it will be created.
        tfdbg_run_id: Debugger Run ID.
        circular_buffer_size: Size of the circular buffer for each of the two
        execution-related debug events files: with the following suffixes: -
        .execution - .graph_execution_traces If <= 0, the circular-buffer
        behavior will be abolished in the constructed object.


def __init__(self, model: PreTrainedModel):
    super().__init__()
    if model.generation_config is None:
        raise AssertionError('The model must have a generation config to be exported with static caching. Please set `generation_config`.')
    if not model.generation_config.use_cache:
        raise AssertionError('The model must have caching enabled to be exported with static caching. Please set `generation_config.use_cache=True`.')
    if model.generation_config.cache_implementation != 'static':
        raise AssertionError("The model must use a 'static' caching implementation to be exported with static caching. Please set `generation_config.cache_implementation='static'`.")
    self.model = model
    self.static_cache = StaticCache(config=self.model.config, max_batch_size=self.model.generation_config.cache_config.batch_size, max_cache_len=self.model.generation_config.cache_config.max_cache_len, device=self.model.generation_config.cache_config.device, dtype=self.model.dtype)
    for i in range(len(self.static_cache.key_cache)):
        self.register_buffer(f'key_cache_{i}', self.static_cache.key_cache[i], persistent=False)
        self.register_buffer(f'value_cache_{i}', self.static_cache.value_cache[i], persistent=False)

Initializes the wrapper module with the pretrained model.

Args:
        model (`PreTrainedModel`): The pretrained model to wrap. The model must have caching
        enabled and use a 'static' caching implementation.

Raises:
        AssertionError: If the pretrained model does not have caching enabled or if it does
        not use a 'static' caching implementation in `model.generation_config`.


def strides(self) -> List[int]:
    return _compute_mesh_strides(self.shape())

Returns the strides tensor array for this mesh.

    If the mesh shape is `[a, b, c, d]`, then the strides array can be computed
    as `[b*c*d, c*d, d, 1]`. This array can be useful in computing local device
    offsets given a device ID. Using the same example, the device coordinates of
    the mesh can be computed as:

    ```
    [(device_id / (b*c*d)) % a,
    (device_id / (c*d))   % b,
    (device_id / (d))     % c,
    (device_id)           % d]
    ```

    This is the same as `(device_id // mesh.strides) % mesh.shape`.

Returns:
        The mesh strides as an integer tensor.


def forward(**super_kwargs):
    super().forward(**super_kwargs)

Example:

    ```python
    >>> import torch
    >>> from transformers import AutoProcessor, AutoModelForImageTextToText

    >>> torch_device = "cuda"
    >>> processor = AutoProcessor.from_pretrained("OpenGVLab/InternVL3-1B-hf")
    >>> model = AutoModelForImageTextToText.from_pretrained(
    ...     "OpenGVLab/InternVL3-1B-hf", torch_dtype=torch.bfloat16, device_map=torch_device
    ... )

    >>> messages = [
    ...     {
    ...         "role": "user",
    ...         "content": [
    ...             {
    ...                 "type": "image",
    ...                 "url": "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg",
    ...             },
    ...             {
    ...                 "type": "image",
    ...                 "url": "https://thumbs.dreamstime.com/b/golden-gate-bridge-san-francisco-purple-flowers-california-echium-candicans-36805947.jpg",
    ...             },
    ...             {"type": "text", "text": "These images depict two different landmarks. Can you identify them?"},
    ...         ],
    ...     },
    ... ]

    >>> inputs = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt").to(torch_device)
    >>> generate_ids = model.generate(**inputs, max_new_tokens=200)
    >>> print(processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True))
    The images depict the Statue of Liberty and the Golden Gate Bridge.
    ```


def sanity_check_tensor_sync(tensor: torch.Tensor, mesh: DeviceMesh, rtol: float=0.0001, atol: float=0.0001, not_sync: bool=False) -> None:
    if not dist.is_initialized() or mesh.size() == 1:
        return
    pg = mesh.get_group()
    if hasattr(tensor, 'to_local'):
        local_tensor = tensor.to_local()
    else:
        local_tensor = tensor
    world_size = dist.get_world_size(pg)
    gathered_tensors = [torch.empty_like(local_tensor) for _ in range(world_size)]
    dist.all_gather(gathered_tensors, local_tensor, group=pg)
    for i in range(1, world_size):
        try:
            torch.testing.assert_close(gathered_tensors[0], gathered_tensors[i], rtol=rtol, atol=atol)
        except AssertionError as e:
            if not_sync:
                continue
            raise e

Verify that a tensor is synchronized (or not synchronized) across all processes in the mesh's process group.
    Handles both regular tensors and DTensors.

Args:
        tensor (torch.Tensor): The tensor to check for synchronization (can be DTensor)
        mesh (DeviceMesh): The device mesh containing the process group
        rtol (float): Relative tolerance for comparison
        atol (float): Absolute tolerance for comparison
        not_sync (bool): If True, asserts that tensors are NOT synchronized. If False, asserts they are synchronized.


def dense_labels_to_sparse(dense, length):
    flat_values = array_ops.reshape(dense, [-1])
    flat_indices = math_ops.range(array_ops.shape(flat_values, out_type=dtypes.int64)[0])
    mask = array_ops.sequence_mask(length, maxlen=array_ops.shape(dense)[1])
    flat_mask = array_ops.reshape(mask, [-1])
    indices = array_ops.expand_dims(array_ops.boolean_mask(flat_indices, flat_mask), 1)
    values = array_ops.boolean_mask(flat_values, flat_mask)
    sparse = sparse_tensor.SparseTensor(indices=indices, values=math_ops.cast(values, dtypes.int32), dense_shape=array_ops.shape(flat_values, out_type=dtypes.int64))
    reshaped = sparse_ops.sparse_reshape(sparse, array_ops.shape(dense))
    max_length = math_ops.reduce_max(length)
    return sparse_tensor.SparseTensor(indices=reshaped.indices, values=reshaped.values, dense_shape=[math_ops.cast(reshaped.dense_shape[0], dtypes.int64), math_ops.cast(max_length, dtypes.int64)])

Convert dense labels with sequence lengths to sparse tensor.

Args:
        dense: tensor of shape [batch, max_length]
        length: int tensor of shape [batch] The length of each sequence in dense.

Returns:
        tf.sparse.SparseTensor with values only for the valid elements of sequences.


def dtype(x):
    return backend.standardize_dtype(x.dtype)

Return the dtype of the tensor input as a standardized string.

    Note that due to the standardization, the dtype will not compare equal
    to the backend-specific version of the dtype.

Args:
        x: A tensor. This function will try to access the `dtype` attribute of
            the input tensor.

Returns:
        A string indicating the dtype of the input tensor, e.g. `"float32"`.

Example:
    >>> x = keras.ops.zeros((8, 12))
    >>> keras.ops.dtype(x)
    'float32'


def __init__(self, max_length: int, eos_token_id: int):
    self.max_length = max_length
    if eos_token_id < 0:
        raise ValueError(f'The forced eos token id must be a non-negative integer, got {eos_token_id}')
    self.eos_token_id = eos_token_id

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

Args:
        max_length (`int`):
            The maximum length of the sequence to be generated.
        eos_token_id (`int`):
            The id of the token to force as the last generated token when `max_length` is reached.


def _MaybeCaptured(t):
    if not isinstance(t, ops.EagerTensor) and _IsFunction(t.op.graph) and (t.op.type == 'Placeholder'):
        for input_t, placeholder_t in _Captures(t.op.graph):
            if t is placeholder_t:
                return _MaybeCaptured(input_t)
    return t

If t is a captured value placeholder, returns the original captured value.

Args:
        t: Tensor

Returns:
        A tensor, potentially from a different Graph/FuncGraph.


def last_updated(self, path):
    return self._gcsIO().last_updated(path)

Get UNIX Epoch time in seconds on the FileSystem.

Args:
        path: string path of file.

    Returns: float UNIX Epoch time

Raises:
        ``BeamIOError``: if path doesn't exist.


def cos(x):
    if any_symbolic_tensors((x,)):
        return Cos().symbolic_call(x)
    return backend.numpy.cos(x)

Cosine, element-wise.

Args:
        x: Input tensor.

Returns:
        The corresponding cosine values.


def make_list_of_t(ts, check_graph=True, allow_graph=True, ignore_ops=False):
    if isinstance(ts, ops.Graph):
        if allow_graph:
            return get_tensors(ts)
        else:
            raise TypeError('allow_graph is False: cannot convert a tf.Graph.')
    else:
        if not is_iterable(ts):
            ts = [ts]
        if not ts:
            return []
        if check_graph:
            check_types = None if ignore_ops else tensor_lib.Tensor
            get_unique_graph(ts, check_types=check_types)
        return [t for t in ts if isinstance(t, tensor_lib.Tensor)]

Convert ts to a list of `tf.Tensor`.

Args:
        ts: can be an iterable of `tf.Tensor`, a `tf.Graph` or a single tensor.
        check_graph: if `True` check if all the tensors belong to the same graph.
        allow_graph: if `False` a `tf.Graph` cannot be converted.
        ignore_ops: if `True`, silently ignore `tf.Operation`.

Returns:
        A newly created list of `tf.Tensor`.

Raises:
        TypeError: if `ts` cannot be converted to a list of `tf.Tensor` or,
        if `check_graph` is `True`, if all the ops do not belong to the same graph.


def __init__(self, keys, values, key_dtype=None, value_dtype=None, name=None):
    if not context.executing_eagerly() and ops.get_default_graph()._get_control_flow_context() is not None:
        with ops.init_scope():
            self._keys = ops.convert_to_tensor(keys, dtype=key_dtype, name='keys')
            self._values = ops.convert_to_tensor(values, dtype=value_dtype, name='values')
    else:
        self._keys = ops.convert_to_tensor(keys, dtype=key_dtype, name='keys')
        self._values = ops.convert_to_tensor(values, dtype=value_dtype, name='values')
    self._name = name if name is not None else 'key_value_init'
    if context.executing_eagerly():
        self._name += str(ops.uid())
    super(KeyValueTensorInitializer, self).__init__(self._keys.dtype, self._values.dtype)

Constructs a table initializer object based on keys and values tensors.

Args:
        keys: The tensor for the keys.
        values: The tensor for the values.
        key_dtype: The `keys` data type. Used when `keys` is a python array.
        value_dtype: The `values` data type. Used when `values` is a python array.
        name: A name for the operation (optional).


def __init__(self, pool_scales: Tuple[int, ...], in_channels: int, channels: int, align_corners: bool) -> None:
    super().__init__()
    self.pool_scales = pool_scales
    self.align_corners = align_corners
    self.in_channels = in_channels
    self.channels = channels
    self.blocks = []
    for i, pool_scale in enumerate(pool_scales):
        block = UperNetPyramidPoolingBlock(pool_scale=pool_scale, in_channels=in_channels, channels=channels)
        self.blocks.append(block)
        self.add_module(str(i), block)

Pyramid Pooling Module (PPM) used in PSPNet.

Args:
        pool_scales (`Tuple[int]`):
            Pooling scales used in Pooling Pyramid Module.
        in_channels (`int`):
            Input channels.
        channels (`int`):
            Channels after modules, before conv_seg.
        align_corners (`bool`):
            align_corners argument of F.interpolate.


def __init__(self, name=None):
    self._name = name
    self._items = []

Menu constructor.

Args:
        name: (str or None) name of this menu.


def __sub__(self, period_tensor):
    return self + periods.PeriodTensor(-period_tensor.quantity(), period_tensor.period_type())

Subtracts a tensor of periods.

    When subtracting months or years, the resulting day of the month is
    decreased to the largest valid value if necessary. E.g. 31.03.2020 - 1 month
    = 29.02.2020, 29.02.2020 - 1 year = 28.02.2019.

Args:
        period_tensor: a PeriodTensor object broadcastable to the shape of "self".

Returns:
        The new instance of DateTensor.


def _get_target_dtype(self, from_dtype: Optional[_NpDType]) -> Optional[_NpDType]:

Validate and normalize the numpy dtype.

Args:
        from_dtype: DType of the array to cast

Returns:
        to_dtype: DType of the array after casting


def try_listify_dict_with_int_keys(src: Dict[Any, Any], convert_when_sparse: bool=False) -> Tuple[Union[List[Any], Dict[Any, Any]], bool]:
    if not src:
        return (src, False)
    min_key = None
    max_key = None
    for key in src.keys():
        if not isinstance(key, int):
            return (src, False)
        if min_key is None or min_key > key:
            min_key = key
        if max_key is None or max_key < key:
            max_key = key
    if convert_when_sparse or (min_key == 0 and max_key == len(src) - 1):
        return ([src[key] for key in sorted(src.keys())], True)
    return (src, False)

Try to convert a dictionary with consequentive integer keys to a list.

Args:
        src: A dict whose keys may be int type and their range form a perfect
        range(0, N) list unless convert_when_sparse is set to True.
        convert_when_sparse: When src is a int-key dict, force convert
        it to a list ordered by key, even it's sparse.

Returns:
        converted list or src unchanged.


def transform_feature(self, transformation_cache, state_manager):
    input_tensor = transformation_cache.get(self.key, state_manager)
    if self.normalizer_fn is not None:
        input_tensor = self.normalizer_fn(input_tensor)
    return input_tensor

See `FeatureColumn` base class.

    In this case, we apply the `normalizer_fn` to the input tensor.

Args:
        transformation_cache: A `FeatureTransformationCache` object to access
        features.
        state_manager: A `StateManager` to create / access resources such as
        lookup tables.

Returns:
        Normalized input tensor.


def __call__(self, shape, dtype=dtypes.float32, **kwargs):
    self._validate_kwargs(kwargs, support_partition=False)
    dtype = _assert_float_dtype(dtype)
    if len(shape) < 2:
        raise ValueError(f'The tensor to initialize, specified by argument `shape` must be at least two-dimensional. Received shape={shape}')
    num_rows = 1
    for dim in shape[:-1]:
        num_rows *= dim
    num_cols = shape[-1]
    flat_shape = (max(num_cols, num_rows), min(num_cols, num_rows))
    a = self._random_generator.random_normal(flat_shape, dtype=dtype)
    q, r = gen_linalg_ops.qr(a, full_matrices=False)
    d = array_ops.diag_part(r)
    q *= math_ops.sign(d)
    if num_rows < num_cols:
        q = array_ops.matrix_transpose(q)
    return self.gain * array_ops.reshape(q, shape)

Returns a tensor object initialized as specified by the initializer.

Args:
        shape: Shape of the tensor.
        dtype: Optional dtype of the tensor. Only floating point types are
        supported.
        **kwargs: Additional keyword arguments.

Raises:
        ValueError: If the dtype is not floating point or the input shape is not
        valid.


def is_null_merge(self):
    return not bool(self._spec.to_string())

Indicate whether the wrapped spec is empty.

    In the degenerate case where self._spec is an empty specification, a caller
    may wish to skip a merge step entirely. (However this class does not have
    enough information to make that determination.)

Returns:
        A boolean indicating whether a device merge will be trivial.


def __init__(self, name='mean_tensor', dtype=None, shape=None):
    super(MeanTensor, self).__init__(name=name, dtype=dtype)
    self._shape = None
    self._total = None
    self._count = None
    self._built = False
    if shape is not None:
        self._build(shape)

Computes the element-wise (weighted) mean of the given tensors.

    `MeanTensor` returns a tensor with the same shape of the input tensors. The
    mean value is updated by keeping local variables `total` and `count`. The
    `total` tracks the sum of the weighted values, and `count` stores the sum of
    the weighted counts.

Args:
        name: (Optional) string name of the metric instance.
        dtype: (Optional) data type of the metric result.
        shape: (Optional) A list of integers, a tuple of integers, or a 1-D Tensor
        of type int32. If not specified, the shape is inferred from the values at
        the first call of update_state.

    Standalone usage:

    >>> m = tf.keras.metrics.MeanTensor()
    >>> m.update_state([0, 1, 2, 3])
    >>> m.update_state([4, 5, 6, 7])
    >>> m.result().numpy()
    array([2., 3., 4., 5.], dtype=float32)

    >>> m.update_state([12, 10, 8, 6], sample_weight= [0, 0.2, 0.5, 1])
    >>> m.result().numpy()
    array([2.       , 3.6363635, 4.8      , 5.3333335], dtype=float32)

    >>> m = tf.keras.metrics.MeanTensor(dtype=tf.float64, shape=(1, 4))
    >>> m.result().numpy()
    array([[0., 0., 0., 0.]])
    >>> m.update_state([[0, 1, 2, 3]])
    >>> m.update_state([[4, 5, 6, 7]])
    >>> m.result().numpy()
    array([[2., 3., 4., 5.]])


def flatten_with_joined_string_paths(structure, separator='/', expand_composites=False):
    flat_paths = yield_flat_paths(structure, expand_composites=expand_composites)

    def stringify_and_join(path_elements):
        return separator.join((str(path_element) for path_element in path_elements))
    flat_string_paths = (stringify_and_join(path) for path in flat_paths)
    return list(zip(flat_string_paths, flatten(structure, expand_composites=expand_composites)))

Returns a list of (string path, atom) tuples.

    The order of tuples produced matches that of `nest.flatten`. This allows you
    to flatten a nested structure while keeping information about where in the
    structure each atom was located. See `nest.yield_flat_paths`
    for more information.

Args:
        structure: the nested structure to flatten.
        separator: string to separate levels of hierarchy in the results, defaults
        to '/'.
        expand_composites: If true, then composite tensors such as
        `tf.sparse.SparseTensor` and `tf.RaggedTensor` are expanded into their
        component tensors.

Returns:
        A list of (string, atom) tuples.


def where(condition, x1=None, x2=None):
    if x1 is None and x2 is not None or (x1 is not None and x2 is None):
        raise ValueError('`x1` and `x2` either both should be `None` or both should have non-None value.')
    if any_symbolic_tensors((condition, x1, x2)):
        return Where().symbolic_call(condition, x1, x2)
    return backend.numpy.where(condition, x1, x2)

Return elements chosen from `x1` or `x2` depending on `condition`.

Args:
        condition: Where `True`, yield `x1`, otherwise yield `x2`.
        x1: Values from which to choose when `condition` is `True`.
        x2: Values from which to choose when `condition` is `False`.

Returns:
        A tensor with elements from `x1` where `condition` is `True`, and
        elements from `x2` where `condition` is `False`.


def _measure_tensor_list_column_widths(self, data):
    max_timestamp_width = 0
    if data:
        max_rel_time_ms = (data[-1].timestamp - self._debug_dump.t0) / 1000.0
        max_timestamp_width = len('[%.3f] ' % max_rel_time_ms) + 1
    max_timestamp_width = max(max_timestamp_width, len(self._TIMESTAMP_COLUMN_HEAD) + 1)
    max_dump_size_width = 0
    for dump in data:
        dump_size_str = cli_shared.bytes_to_readable_str(dump.dump_size_bytes)
        if len(dump_size_str) + 1 > max_dump_size_width:
            max_dump_size_width = len(dump_size_str) + 1
    max_dump_size_width = max(max_dump_size_width, len(self._DUMP_SIZE_COLUMN_HEAD) + 1)
    max_op_type_width = 0
    for dump in data:
        op_type = self._debug_dump.node_op_type(dump.node_name)
        if len(op_type) + 1 > max_op_type_width:
            max_op_type_width = len(op_type) + 1
    max_op_type_width = max(max_op_type_width, len(self._OP_TYPE_COLUMN_HEAD) + 1)
    return (max_timestamp_width, max_dump_size_width, max_op_type_width)

Determine the maximum widths of the timestamp and op-type column.

    This method assumes that data is sorted in the default order, i.e.,
    by ascending timestamps.

Args:
        data: (list of DebugTensorDaum) the data based on which the maximum
        column widths will be determined.

Returns:
        (int) maximum width of the timestamp column. 0 if data is empty.
        (int) maximum width of the dump size column. 0 if data is empty.
        (int) maximum width of the op type column. 0 if data is empty.


def broadcast(self, tensor, destinations):
    validate_destinations(destinations)
    return self.broadcast_implementation(tensor, destinations)

Broadcast `tensor` to `destinations`.

    This can only be called in the cross-replica context.

Args:
        tensor: a `tf.Tensor` like object. The value to broadcast.
        destinations: a `tf.distribute.DistributedValues`, a `tf.Variable`, a
        `tf.Tensor` alike object, or a device string. It specifies the devices
        to broadcast to. Note that if it's a `tf.Variable`, the value is
        broadcasted to the devices of that variable, this method doesn't update
        the variable.

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


def __init__(self, losses):
    super().__init__()
    self.loss_map = {'iou': self.loss_iou, 'distance': self.loss_distance, 'duration': self.loss_duration}
    for loss in losses:
        if loss not in self.loss_map:
            raise ValueError(f'Loss {loss} not supported')
    self.losses = losses

This class computes the losses for `TvpForVideoGrounding`. The process happens in two steps: 1) we compute
    hungarian assignment between ground truth boxes and the outputs of the model 2) we supervise each pair of matched
    ground-truth / prediction (supervise class and box).

Args:
        losses (`List[str]`):
            List of all the losses to be applied.


Test function

Args:
        x: The input


def compute_optimal_blocks(device: torch.device, config: PretrainedConfig, generation_config: GenerationConfig, inputs: List[List[int]], dtype: torch.dtype=torch.bfloat16, safety_margin: float=0.9, median_prefill_length: Optional[int]=None):
    head_dim = getattr(config, 'head_dim', config.hidden_size // config.num_attention_heads)
    num_kv_heads = getattr(config, 'num_key_value_heads', config.num_attention_heads)
    num_hidden_layers = getattr(config, 'num_hidden_layers', 40)
    if device.type == 'cuda':
        device_properties = torch.cuda.get_device_properties(device)
        total_memory = device_properties.total_memory
        allocated_memory = torch.cuda.memory_allocated(device)
        reserved_memory = torch.cuda.memory_reserved(device)
        available_memory = total_memory - max(allocated_memory, reserved_memory)
    elif device.type == 'mps':
        logger.warning('MPS memory estimation is approximate. Using conservative defaults.')
        return (2048, 256)
    else:
        logger.warning(f'Unsupported device type {device.type} for optimal block calculation. Using defaults.')
        return (32, 128)
    available_memory = int(available_memory * safety_margin)
    if available_memory <= 0:
        logger.warning('Not enough available memory. Using minimum configuration.')
        return (8, 128)
    dtype_size = torch.tensor([], dtype=dtype).element_size()
    memory_per_token = 2 * num_kv_heads * head_dim * dtype_size * num_hidden_layers
    tokens_to_generate = getattr(generation_config, 'max_new_tokens', 20)
    if median_prefill_length is None and inputs:
        non_empty_inputs = [len(seq) for seq in inputs if seq]
        median_prefill_length = int(statistics.median(non_empty_inputs)) if non_empty_inputs else 64
    elif median_prefill_length is None:
        median_prefill_length = 64
    seq_length = median_prefill_length + tokens_to_generate
    MIN_BLOCK_SIZE = 16
    per_sequence_memory = seq_length * memory_per_token
    max_concurrent_sequences = max(1, int(available_memory // per_sequence_memory))
    total_tokens = available_memory // memory_per_token
    initial_block_size = max(MIN_BLOCK_SIZE, total_tokens // (max_concurrent_sequences * 2))
    block_size = 1 << (initial_block_size - 1).bit_length()
    num_blocks = max(1, total_tokens // block_size)
    logger.info(f'Optimal cache: {num_blocks} blocks of size {block_size} (can handle ~{num_blocks * block_size // seq_length} sequences of length {seq_length})')
    return (int(num_blocks), int(block_size))

Calculate optimal number and size of blocks for the KV cache.

Args:
        device: The device where the model runs
        config: The model configuration
        generation_config: The generation configuration
        inputs: Sample input sequences to estimate memory requirements
        dtype: Data type for cache tensors
        safety_margin: Fraction of available memory to use
        median_prefill_length: Override for median prefill length calculation

Returns:
        Tuple of (num_blocks, block_size)


def price(self, valuation_date, market, model=None, pricing_context=None, name=None):
    del model, pricing_context
    name = name or self._name + '_price'
    with tf.name_scope(name):
        discount_curve = market.discount_curve
        reference_curve = market.reference_curve
        libor_rate = rc.get_rate_index(market, self._start_date, rc.RateIndexType.LIBOR, dtype=self._dtype)
        libor_rate = tf.repeat(tf.convert_to_tensor(libor_rate, dtype=self._dtype), self._num_cashflows)
        discount_factors = discount_curve.get_discount_factor(self._payment_dates)
        forward_rates = reference_curve.get_forward_rate(self._accrual_start_date, self._accrual_end_date, self._daycount_fractions)
        forward_rates = tf.where(self._daycount_fractions > 0.0, forward_rates, tf.zeros_like(forward_rates))
        forward_rates = tf.where(self._coupon_end_dates < valuation_date, tf.constant(0.0, dtype=self._dtype), tf.where(self._coupon_start_dates < valuation_date, libor_rate, forward_rates))
        coupon_rate = self._coupon_multiplier * (forward_rates + self._coupon_basis)
        cashflow_pvs = self._notional * (self._daycount_fractions * coupon_rate * discount_factors)
        return tf.math.reduce_sum(tf.reshape(cashflow_pvs, (self._batch_size, self._num_cashflows)), axis=1)

Returns the present value of the stream on the valuation date.

Args:
        valuation_date: A scalar `DateTensor` specifying the date on which
        valuation is being desired.
        market: A namedtuple of type `InterestRateMarket` which contains the
        necessary information for pricing the cashflow stream.
        model: Reserved for future use.
        pricing_context: Additional context relevant for pricing.
        name: Python str. The name to give to the ops created by this function.
        Default value: `None` which maps to 'price'.

Returns:
        A Rank 1 `Tensor` of real type containing the modeled price of each stream
        contract based on the input market data.


def __call__(self, shape, dtype=None):
    dtype = standardize_dtype(dtype)
    return ops.zeros(shape, dtype=dtype)

Returns a tensor object initialized as specified by the initializer.

Args:
        shape: Shape of the tensor.
        dtype: Optional dtype of the tensor. Only numeric or boolean dtypes
            are supported. If not specified, `keras.backend.floatx()`
            is used, which default to `float32` unless you configured it
            otherwise (via `keras.backend.set_floatx(float_dtype)`).


def assign(self, value, use_locking=None, name=None, read_value=True):
    with _handle_graph(self.handle):
        value_tensor = ops.convert_to_tensor(value, dtype=self.dtype)
        if not self._shape.is_compatible_with(value_tensor.shape):
            if self.name is None:
                tensor_name = ''
            else:
                tensor_name = ' ' + str(self.name)
            raise ValueError(f"Cannot assign value to variable '{tensor_name}': Shape mismatch.The variable shape {self._shape}, and the assigned value shape {value_tensor.shape} are incompatible.")
        kwargs = {}
        if forward_compat.forward_compatible(2022, 3, 23):
            validate_shape = self._validate_shape and self._shape.is_fully_defined()
            kwargs['validate_shape'] = validate_shape
        assign_op = gen_resource_variable_ops.assign_variable_op(self.handle, value_tensor, name=name, **kwargs)
        if read_value:
            return self._lazy_read(assign_op)
    return assign_op

Assigns a new value to this variable.

Args:
        value: A `Tensor`. The new value for this variable.
        use_locking: If `True`, use locking during the assignment.
        name: The name to use for the assignment.
        read_value: A `bool`. Whether to read and return the new value of the
        variable or not.

Returns:
        If `read_value` is `True`, this method will return the new value of the
        variable after the assignment has completed. Otherwise, when in graph mode
        it will return the `Operation` that does the assignment, and when in eager
        mode it will return `None`.


def GetFile(self, map_name, dst_file, current_file, location=None):
    if map_name == config.MAP_PASSWORD:
        return self.GetPasswdFile(dst_file, current_file)
    elif map_name == config.MAP_GROUP:
        return self.GetGroupFile(dst_file, current_file)
    elif map_name == config.MAP_SHADOW:
        return self.GetShadowFile(dst_file, current_file)
    elif map_name == config.MAP_NETGROUP:
        return self.GetNetgroupFile(dst_file, current_file)
    elif map_name == config.MAP_AUTOMOUNT:
        return self.GetAutomountFile(dst_file, current_file, location=location)
    raise error.UnsupportedMap('Source can not fetch %s' % map_name)

Retrieve a file from this source.

Args:
        map_name: A string representation of the map whose file you want
        dst_file: Temporary filename to write to.
        current_file: Path to the current cache.
        location: optional field used by automounts to indicate a specific map

Returns:
        path to new file

Raises:
        UnsupportedMap: for unknown source maps


def lookup(self, keys, name=None):
    if keys.dtype.base_dtype != self._key_dtype:
        raise TypeError(f'Dtype of argument `keys` must be {self._key_dtype}, received: {keys.dtype}')
    values = keys
    if isinstance(keys, (sparse_tensor.SparseTensor, internal.RaggedTensor)):
        values = keys.values
    if self._table and self._table.key_dtype.base_dtype == dtypes.int64:
        values = math_ops.cast(values, dtypes.int64)
    if self._num_oov_buckets == 0:
        ids = self._table.lookup(values, name=name)
    else:
        with ops.name_scope(name, '%s_Lookup' % self.name):
            str_to_hash_bucket = self._get_string_to_hash_bucket_fn(self._hasher_spec)
            buckets = str_to_hash_bucket(_as_string(values), num_buckets=self._num_oov_buckets, name='hash_bucket')
            if self._table:
                ids = self._table.lookup(values)
                buckets = math_ops.add(buckets, self._table.size())
                is_id_non_default = math_ops.not_equal(ids, self._table.default_value)
                ids = array_ops.where_v2(is_id_non_default, ids, buckets)
            else:
                ids = buckets
    if isinstance(keys, sparse_tensor.SparseTensor):
        return sparse_tensor.SparseTensor(keys.indices, ids, keys.dense_shape)
    elif isinstance(keys, internal.RaggedTensor):
        return keys.with_values(ids)
    return ids

Looks up `keys` in the table, outputs the corresponding values.

    It assigns out-of-vocabulary keys to buckets based in their hashes.

Args:
        keys: Keys to look up. May be either a `SparseTensor` or dense `Tensor`.
        name: Optional name for the op.

Returns:
        A `SparseTensor` if keys are sparse, a `RaggedTensor` if keys are ragged,
        otherwise a dense `Tensor`.

Raises:
        TypeError: when `keys` doesn't match the table key data type.


def report_run(config, auth, account, report_id=None, name=None):
    account_id, advertiser_id = parse_account(config, auth, account)
    is_superuser, profile_id = get_profile_for_api(config, auth, account_id)
    kwargs = {'profileId': profile_id, 'accountId': account_id} if is_superuser else {'profileId': profile_id}
    if config.verbose:
        print('DCM REPORT RUN INIT', report_id or name)
    if report_id is None:
        report = report_get(config, auth, account, name=name)
        if report is None:
            raise Exception('Report does not exist:', name)
        else:
            report_id = report['id']
    kwargs = {'profileId': profile_id, 'accountId': account_id} if is_superuser else {'profileId': profile_id}
    kwargs['reportId'] = report_id
    files = report_files(config, auth, account, report_id)
    latest_file_json = next(files, None)
    if latest_file_json == None or latest_file_json['status'] != 'PROCESSING':
        if config.verbose:
            print('RUNNING REPORT', report_id or name)
        API_DCM(config, auth, internal=is_superuser).reports().run(**kwargs).execute()
        return True
    if config.verbose:
        print('REPORT RUN SKIPPED', report_id or name)
    return False

Trigger a DCM report to run by name or ID.

    Will do nothing if report is currently in progress.

    Bulletproofing:
    https://developers.google.com/doubleclick-advertisers/v3.3/reports/run

Args:
        * auth: (string) Either user or service.
        * account: (string) [account:advertiser@profile] token.
        * report_id: (int) ID of DCm report to fetch ( either or name ).
        * name: (string) Name of report to fetch ( either or report_id ).

Returns:
        * True if report run is executed
        * False otherwise


def send_graph_tracebacks(destinations, run_key, origin_stack, graph, send_source=True):
    _send_call_tracebacks(destinations, origin_stack, is_eager_execution=False, call_key=run_key, graph=graph, send_source=send_source)

Send the tracebacks of a graph execution call to debug server(s).

Args:
        destinations: gRPC destination addresses, a `str` or a `list` of `str`s,
        e.g., "localhost:4242". If a `list`, gRPC requests containing the same
        `CallTraceback` proto payload will be sent to all the destinations.
        run_key: A string describing the feeds, fetches (and targets) names of the
        `tf.Session.run` call.
        origin_stack: The traceback of the `tf.Session.run()` invocation.
        graph: A Python `tf.Graph` object (i.e., *not* a `tf.compat.v1.GraphDef`),
        which contains op tracebacks.
        send_source: Whether the source files involved in the op tracebacks but
        outside the TensorFlow library are to be sent.


def get_creds(use_personal_account: bool, service_account: str, private_key: str) -> credentials.Credentials:
    if service_account and private_key:
        try:
            with open_local(private_key) as local_path:
                creds = ee.ServiceAccountCredentials(service_account, local_path)
        except Exception:
            raise RuntimeError(f'Unable to open the private key {private_key}.')
    elif use_personal_account:
        ee.Authenticate()
        creds, _ = default()
    elif is_compute_engine():
        creds = compute_engine.Credentials()
    else:
        creds, _ = default()
    creds.refresh(requests.Request())
    return creds

Fetches credentials for authentication.

    If the `use_personal_account` argument is true then it will authenticate with pop-up
    browser window using personal account. Otherwise, if the application is running
    in compute engine, it will use credentials of service account bound to the VM.
    Otherwise, it will try to use user credentials.

Args:
        use_personal_account: A flag to use personal account for ee authentication.
        service_account: Service account address when using a private key for earth engine authentication.
        private_key: A private key path to authenticate earth engine using private key.

Returns:
        cred: Credentials object.


def dismantle_graph(graph) -> None:
    graph._functions.clear()
    graph.Dismantle()

Cleans up reference cycles from a `Graph`.

    Helpful for making sure the garbage collector doesn't need to run after a
    temporary `Graph` is no longer needed.

Args:
        graph: A `Graph` object to destroy. Neither it nor any of its ops are usable
        after this function runs.


def get_metrics(pred: jax.Array, actual: jax.Array) -> Result:
    tp: int = jnp.sum(jnp.logical_and(pred == 1, actual == 1))
    tn: int = jnp.sum(jnp.logical_and(pred == 0, actual == 0))
    fp: int = jnp.sum(jnp.logical_and(pred == 1, actual == 0))
    fn: int = jnp.sum(jnp.logical_and(pred == 0, actual == 1))
    accuracy = (tp + tn) / (tp + tn + fp + fn + EPS)
    precision = tp / (tp + fp + EPS)
    recall = tp / (tp + fn + EPS)
    fscore = 2 * precision * recall / (precision + recall + EPS)
    return Result(tp=tp, tn=tn, fp=fp, fn=fn, accuracy=accuracy, precision=precision, recall=recall, fscore=fscore)

Gets evaluation metrics from the prediction and the actual target.

Args:
        pred (jax.Array): A prediction of the target.
        actual (jax.Array): The actual target.

Returns:
        result (Result): A result.


def _gather_saveables_for_checkpoint(self) -> Dict[str, Callable[..., Any]]:

    def _saveable_factory(name=self._common_name):
        saveables = []
        num_shards = len(self.values)
        for shard_id in range(num_shards):
            saveables.append(TPUEmbeddingShardedSaveable(self.values[shard_id], shard_id, num_shards, self.shard_dim, name))
        return saveables
    return {base.VARIABLE_VALUE_KEY: _saveable_factory}

Overrides Trackable method.

Returns:
        A dictionary mapping attribute names to `SaveableObject` factories.


def get_image_tokens(self, pixel_values: torch.FloatTensor):
    batch_size = pixel_values.shape[0]
    _, _, image_toks = self.vqmodel.encode(pixel_values)
    bpe_toks = self.vocabulary_mapping.convert_img2bpe(image_toks)
    bpe_toks = bpe_toks.view(batch_size, -1)
    return bpe_toks

Tokenizes images into discrete tokens with VQGAN module. Converts
    obtained image tokens into BPE tokens and wraps with "boi" and "eoi"
    special tokens.

Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)):
            The tensors corresponding to the input images.


def get_collection(key, scope=None) -> list[Any]:
    return get_default_graph().get_collection(key, scope)

Wrapper for `Graph.get_collection()` using the default graph.

    See `tf.Graph.get_collection`
    for more details.

Args:
        key: The key for the collection. For example, the `GraphKeys` class contains
        many standard names for collections.
        scope: (Optional.) If supplied, the resulting list is filtered to include
        only items whose `name` attribute matches using `re.match`. Items without
        a `name` attribute are never returned if a scope is supplied and the
        choice or `re.match` means that a `scope` without special tokens filters
        by prefix.

Returns:
        The list of values in the collection with the given `name`, or
        an empty list if no value has been added to that collection. The
        list contains the values in the order under which they were
        collected.

    @compatibility(eager)
    Collections are not supported when eager execution is enabled.
    @end_compatibility


def _get_value_type(cls, value):
    type_ = cls.typeDict.get(type(value))
    if type_ is None:
        type_ = 'CLASS' if inspect.isclass(value) else None
    if type_ is None and value is None:
        type_ = 'STRING'
    return type_

Infers the type of a given value.

Args:
        value: The value whose type needs to be inferred. For 'DURATION' and
        'TIMESTAMP', the corresponding Python type is datetime.timedelta and
        datetime.datetime respectively. For Python classes, the API type is
        just 'STRING' at the moment.

Returns:
        One of 'STRING', 'INTEGER', 'FLOAT', 'CLASS', 'DURATION', or
        'TIMESTAMP', depending on the type of the value.


def _add_train_op(self, train_op):
    if train_op is not None:
        if not isinstance(train_op, tensor.Tensor) and (not isinstance(train_op, ops.Operation)):
            raise TypeError(f'`train_op` {train_op} needs to be a Tensor or Op.')
        ops.add_to_collection(constants.TRAIN_OP_KEY, train_op)

Add train op to the SavedModel.

    Note that this functionality is in development, and liable to be
    moved elsewhere.

Args:
        train_op: Op or group of ops that are used for training. These are stored
        as a collection with key TRAIN_OP_KEY, but not executed.

Raises:
        TypeError if Train op is not of type `Operation`.


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

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

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

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

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


def __getitem__(self, indices):
    return self.array[indices]

Select elements in the 0th dimension.

Args:
        indices: the indices to select. Only needs to support one dimension,
            the 0th dimension. Should support a `slice` or a list, tuple,
            `np.array` or 1D tensor.
    Returns: A slice of `self.array`.


def _assert_sample_values_all_close(self, sess: session.Session, repr_ds_1: repr_dataset.RepresentativeDataset, repr_ds_2: repr_dataset.RepresentativeDataset) -> None:
    for sample_1, sample_2 in zip(repr_ds_1, repr_ds_2):
        self.assertCountEqual(sample_1.keys(), sample_2.keys())
        for input_key in sample_1:
            self._assert_tensorlike_all_close(sess, sample_1[input_key], sample_2[input_key])

Asserts that the sample values are "all close" between the two datasets.

    This assumes that the order of corresponding samples is preserved and the
    size of the two datasets are equal.

Args:
        sess: Session instance used to evaluate any tf.Tensors.
        repr_ds_1: A RepresentativeDataset.
        repr_ds_2: A RepresentativeDataset.


def better_logging():
    app.call_after_init(_better_logging)

Improve Python logging when running locally.

    * Display Python logs by default (even when user forgot `--logtostderr`),
        without being polluted by hundreds of C++ logs.
    * Cleaner minimal log format (e.g. `I 15:04:05 [main.py:24]:`)
    * Avoid visual artifacts between TQDM & `logging`
    * Clickable hyperlinks redirecting to code search (require terminal support)

    Usage:

    ```python
    if __name__ == '__main__':
        eapp.better_logging()
        app.run(main)
    ```

    Note this has only effect when user run locally and without `--logtostderr`.


def create_database_view(self, view: views.View, view_name: str) -> None:
    view_sql = f'CREATE OR REPLACE VIEW {self._view_dataset}.{view_name} AS\n{self.to_sql(view)}'
    self._engine.execute(view_sql).fetchall()

Creates a Spark view with the given name in the runner's view_dataset.

Args:
        view: the FHIR view that creates
        view_name: the view name passed to the CREATE OR REPLACE VIEW statement.


def distribution(namespace: Union[Type, str], name: str) -> 'Metrics.DelegatingDistribution':
    namespace = Metrics.get_namespace(namespace)
    return Metrics.DelegatingDistribution(MetricName(namespace, name))

Obtains or creates a Distribution metric.

    Distribution metrics are restricted to integer-only distributions.

Args:
        namespace: A class or string that gives the namespace to a metric
        name: A string that gives a unique name to a metric

Returns:
        A Distribution object.


def join(path, *paths):
    path_ = compat.as_str_any(compat.path_to_str(path))
    if '://' in path_[1:]:
        return urljoin(path, *paths)
    return os.path.join(path, *paths)

Join one or more path components intelligently.

    TensorFlow specific filesystems will be joined
    like a url (using "/" as the path seperator) on all platforms:

    On Windows or Linux/Unix-like:
    >>> tf.io.gfile.join("gcs://folder", "file.py")
    'gcs://folder/file.py'

    >>> tf.io.gfile.join("ram://folder", "file.py")
    'ram://folder/file.py'

    But the native filesystem is handled just like os.path.join:

    >>> path = tf.io.gfile.join("folder", "file.py")
    >>> if os.name == "nt":
    ...   expected = "folder\\file.py"  # Windows
    ... else:
    ...   expected = "folder/file.py"  # Linux/Unix-like
    >>> path == expected
    True

Args:
        path: string, path to a directory
        paths: string, additional paths to concatenate

Returns:
        path: the joined path.


def convert_nested_bidirectional(weights):
    num_weights_per_layer = len(weights) // 2
    forward_weights = preprocess_weights_for_loading(layer.forward_layer, weights[:num_weights_per_layer], original_keras_version, original_backend)
    backward_weights = preprocess_weights_for_loading(layer.backward_layer, weights[num_weights_per_layer:], original_keras_version, original_backend)
    return forward_weights + backward_weights

Converts layers nested in `Bidirectional` wrapper.

    This function uses `preprocess_weights_for_loading()` for converting
    layers.

Args:
        weights: List of weights values (Numpy arrays).

Returns:
        A list of weights values (Numpy arrays).


def __init__(self, graph_parents=None, is_constant_jacobian=False, validate_args=False, dtype=None, forward_min_event_ndims=None, inverse_min_event_ndims=None, name=None):
    self._graph_parents = graph_parents or []
    if forward_min_event_ndims is None and inverse_min_event_ndims is None:
        raise ValueError('Must specify at least one of `forward_min_event_ndims` and `inverse_min_event_ndims`.')
    elif inverse_min_event_ndims is None:
        inverse_min_event_ndims = forward_min_event_ndims
    elif forward_min_event_ndims is None:
        forward_min_event_ndims = inverse_min_event_ndims
    if not isinstance(forward_min_event_ndims, int):
        raise TypeError('Expected forward_min_event_ndims to be of type int, got {}'.format(type(forward_min_event_ndims).__name__))
    if not isinstance(inverse_min_event_ndims, int):
        raise TypeError('Expected inverse_min_event_ndims to be of type int, got {}'.format(type(inverse_min_event_ndims).__name__))
    if forward_min_event_ndims < 0:
        raise ValueError('forward_min_event_ndims must be a non-negative integer.')
    if inverse_min_event_ndims < 0:
        raise ValueError('inverse_min_event_ndims must be a non-negative integer.')
    self._forward_min_event_ndims = forward_min_event_ndims
    self._inverse_min_event_ndims = inverse_min_event_ndims
    self._is_constant_jacobian = is_constant_jacobian
    self._constant_ildj_map = {}
    self._validate_args = validate_args
    self._dtype = dtype
    self._from_y = {}
    self._from_x = {}
    if name:
        self._name = name
    else:

        def camel_to_snake(name):
            s1 = re.sub('(.)([A-Z][a-z]+)', '\\1_\\2', name)
            return re.sub('([a-z0-9])([A-Z])', '\\1_\\2', s1).lower()
        self._name = camel_to_snake(type(self).__name__.lstrip('_'))
    for i, t in enumerate(self._graph_parents):
        if t is None or not tensor_util.is_tf_type(t):
            raise ValueError('Graph parent item %d is not a Tensor; %s.' % (i, t))

Constructs Bijector.

    A `Bijector` transforms random variables into new random variables.

Example:
    ```python
    # Create the Y = g(X) = X transform.
    identity = Identity()

    # Create the Y = g(X) = exp(X) transform.
    exp = Exp()
    ```

    See `Bijector` subclass docstring for more details and specific examples.

Args:
        graph_parents: Python list of graph prerequisites of this `Bijector`.
        is_constant_jacobian: Python `bool` indicating that the Jacobian matrix is
        not a function of the input.
        validate_args: Python `bool`, default `False`. Whether to validate input
        with asserts. If `validate_args` is `False`, and the inputs are invalid,
        correct behavior is not guaranteed.
        dtype: `tf.dtype` supported by this `Bijector`. `None` means dtype is not
        enforced.
        forward_min_event_ndims: Python `integer` indicating the minimum number of
        dimensions `forward` operates on.
        inverse_min_event_ndims: Python `integer` indicating the minimum number of
        dimensions `inverse` operates on. Will be set to
        `forward_min_event_ndims` by default, if no value is provided.
        name: The name to give Ops created by the initializer.

Raises:
        ValueError:  If neither `forward_min_event_ndims` and
        `inverse_min_event_ndims` are specified, or if either of them is
        negative.
        ValueError:  If a member of `graph_parents` is not a `Tensor`.


def auto_shard_dataset(dataset, num_shards, index, num_replicas_in_sync=None):
    if isinstance(dataset, distribute_types.DistributedDatasetInterface):
        return dataset.auto_shard(num_shards, index)
    if dataset.options().experimental_distribute.auto_shard_policy != AutoShardPolicy.OFF:
        if num_replicas_in_sync is None:
            num_replicas_in_sync = 1
        if isinstance(dataset, data_types.DatasetV1):
            return distribute._AutoShardDatasetV1(dataset, num_shards, index, num_replicas_in_sync)
        else:
            return distribute._AutoShardDataset(dataset, num_shards, index, num_replicas_in_sync)
    else:
        return dataset

Shard the input pipeline by sharding the underlying list of files.

Args:
        dataset: A `tf.data.Dataset` instance, typically the result of a bunch of
        dataset transformations.
        num_shards: A `tf.int64` scalar `tf.Tensor`, representing the number of
        shards operating in parallel. Same usage as in `tf.data.Dataset.shard`.
        index: A `tf.int64` scalar `tf.Tensor`, representing the worker index.
        Same usage as in `tf.data.Dataset.shard`.
        num_replicas_in_sync: An integer representing the total number of replicas
        across all workers. This is used in the rewrite when sharding by data.

Returns:
        A modified `Dataset` obtained by updating the pipeline sharded by the
        files. The input dataset will be returned if we cannot automatically
        determine a good way to shard the input dataset.


def index_add(x, idx, y):
    return _index_update_helper(tf_np.ndarray._with_index_add, x, idx, y)

Pure equivalent of `x[idx] += y`.

    Returns the value of x that would result from the NumPy-style indexed
    assignment `x[idx] += y`. Because it's a pure function, `x` itself won't be
    changed.

Args:
        x: an array with the values to be updated.
        idx: a Numpy-style index, consisting of `None`, integers, slice objects,
        ellipses, ndarrays with integer dtypes, or a tuple of the above.
        y: the array of updates. `y` must be broadcastable to the shape of the array
        that would be returned by `x[idx]`.

Returns:
        The updated version of `x`.


def rename(self, source_file_names, destination_file_names):
    err_msg = 'source_file_names and destination_file_names should be equal in length'
    assert len(source_file_names) == len(destination_file_names), err_msg

    def _rename_file(source, destination):
        """Rename a single file object"""
        try:
            os.rename(source, destination)
        except OSError as err:
            raise IOError(err)
    exceptions = {}
    for source, destination in zip(source_file_names, destination_file_names):
        try:
            _rename_file(source, destination)
        except Exception as e:
            exceptions[source, destination] = e
    if exceptions:
        raise BeamIOError('Rename operation failed', exceptions)

Rename the files at the source list to the destination list.
    Source and destination lists should be of the same size.

Args:
        source_file_names: List of file paths that need to be moved
        destination_file_names: List of destination_file_names for the files

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


def shape(input, name=None, out_type=None):
    if out_type is None:
        if flags.config().tf_shape_default_int64.value():
            out_type = dtypes.int64
        else:
            out_type = dtypes.int32
    return shape_internal(input, name, optimize=True, out_type=out_type)

Returns the shape of a tensor.

    This operation returns a 1-D integer tensor representing the shape of `input`.

    For example:

    ```python
    t = tf.constant([[[1, 1, 1], [2, 2, 2]], [[3, 3, 3], [4, 4, 4]]])
    tf.shape(t)  # [2, 2, 3]
    ```

Args:
        input: A `Tensor` or `SparseTensor`.
        name: A name for the operation (optional).
        out_type: (Optional) The specified output type of the operation (`int32`
        or `int64`). Defaults to `tf.int32`.

Returns:
        A `Tensor` of type `out_type`.


def _powerset(iterable):
    s = list(iterable)
    return itertools.chain.from_iterable((itertools.combinations(s, r) for r in range(len(s) + 1)))

Helper for generating all possible reduction_axes arguments.

    Example: powerset([0,1,2]): () (0,) (1,) (2,) (0,1) (0,2) (1,2) (0,1,2)

Args:
        iterable: An iterable of items to generate the powerset of.

Returns:
        The powerset of all items in iterable.


def max_steps_per_epoch():
    return _MAX_STEPS_PER_EPOCH

Get the maximum number of steps for any call to fit/evaluate/predict.

    Retrieves the limit on the number of epochs set by
    `keras.config.set_max_steps_per_epoch` or the `KERAS_MAX_STEPS_PER_EPOCH`
    environment variable.

Args:
        max_epochs: The integer limit on the number of epochs or `None`. If
            `None`, no limit is applied.


def _apply_scores(self, scores, value, scores_mask=None, training=False):
    if scores_mask is not None:
        padding_mask = ops.logical_not(scores_mask)
        max_value = 65504.0 if scores.dtype == 'float16' else 1000000000.0
        scores -= max_value * ops.cast(padding_mask, dtype=scores.dtype)
    weights = ops.softmax(scores, axis=-1)
    if training and self.dropout > 0:
        weights = backend.random.dropout(weights, self.dropout, seed=self.seed_generator)
    return (ops.matmul(weights, value), weights)

Applies attention scores to the given value tensor.

    To use this method in your attention layer, follow the steps:

    * Use `query` tensor of shape `(batch_size, Tq)` and `key` tensor of
        shape `(batch_size, Tv)` to calculate the attention `scores`.
    * Pass `scores` and `value` tensors to this method. The method applies
        `scores_mask`, calculates
        `attention_distribution = softmax(scores)`, then returns
        `matmul(attention_distribution, value).
    * Apply `query_mask` and return the result.

Args:
        scores: Scores float tensor of shape `(batch_size, Tq, Tv)`.
        value: Value tensor of shape `(batch_size, Tv, dim)`.
        scores_mask: A boolean mask tensor of shape `(batch_size, 1, Tv)`
            or `(batch_size, Tq, Tv)`. If given, scores at positions where
            `scores_mask==False` do not contribute to the result. It must
            contain at least one `True` value in each line along the last
            dimension.
        training: Python boolean indicating whether the layer should behave
            in training mode (adding dropout) or in inference mode
            (no dropout).

Returns:
        Tensor of shape `(batch_size, Tq, dim)`.
        Attention scores after masking and softmax with shape
            `(batch_size, Tq, Tv)`.


def _generate_pickle_name(gt):
    grammar_textfile_name = os.path.basename(gt)
    head, tail = os.path.splitext(grammar_textfile_name)
    if tail == '.txt':
        tail = ''
    cache_dir = user_cache_dir(appname='YAPF', appauthor='Google', version=yapf_version)
    return cache_dir + os.sep + head + tail + '-py' + '.'.join(map(str, sys.version_info)) + '.pickle'

Get the filepath to write a pickle file to
    given the path of a grammar textfile.

    The returned filepath should be in a user-specific cache directory.

Args:
        gt (str): path to grammar text file

Returns:
        str: path to pickle file


def _CheckGrayscaleImage(image, require_static=True):
    try:
        if image.get_shape().ndims is None:
            image_shape = image.get_shape().with_rank(2)
        else:
            image_shape = image.get_shape().with_rank_at_least(2)
    except ValueError:
        raise ValueError('A grayscale image (shape %s) must be at least two-dimensional.' % image.shape)
    if require_static and (not image_shape.is_fully_defined()):
        raise ValueError("'image' must be fully defined.")
    if image_shape.is_fully_defined():
        if image_shape[-1] != 1:
            raise ValueError('Last dimension of a grayscale image should be size 1.')
    if not image_shape.is_fully_defined():
        return [check_ops.assert_equal(array_ops.shape(image)[-1], 1, message='Last dimension of a grayscale image should be size 1.'), check_ops.assert_greater_equal(array_ops.rank(image), 3, message='A grayscale image must be at least two-dimensional.')]
    else:
        return []

Assert that we are working with properly shaped grayscale image.

Args:
        image: >= 2-D Tensor of size [*, 1]
        require_static: Boolean, whether static shape is required.

Raises:
        ValueError: if image.shape is not a [>= 2] vector or if
                last dimension is not size 1.

Returns:
        An empty list, if `image` has fully defined dimensions. Otherwise, a list
        containing an assert op is returned.


def convert_to_rgb(video: np.array, data_format: Optional[ChannelDimension]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> np.array:
    if not isinstance(video, np.ndarray):
        raise ValueError(f'Video has to be a numpy array to convert to RGB format, but found {type(video)}')
    if input_data_format is None:
        input_data_format = infer_channel_dimension_format(video)
    video = to_channel_dimension_format(video, ChannelDimension.FIRST, input_channel_dim=input_data_format)
    if video.shape[-3] == 3:
        return video
    if video.shape[-3] == 1:
        return video.repeat(3, -3)
    if not (video[..., 3, :, :] < 255).any():
        return video
    alpha = video[..., 3, :, :] / 255.0
    video = (1 - alpha[..., None, :, :]) * 255 + alpha[..., None, :, :] * video[..., 3, :, :]
    return video

Convert video to RGB by blending the transparency layer if it's in RGBA format, otherwise simply returns it.

Args:
        video (`np.array`):
            The video to convert.
        data_format (`ChannelDimension`, *optional*):
            The channel dimension format of the output video. If unset, will use the inferred format from the input.
        input_data_format (`ChannelDimension`, *optional*):
            The channel dimension format of the input video. If unset, will use the inferred format from the input.


def parallel_part_functions(fns: Sequence[PartFn], match_fns: Sequence[MatchFn] | None=None, with_default_output: bool=False, with_always_output: bool=False) -> PartFn:
    return functools.partial(_parallel_part_functions, _to_tuple_fns(fns, match_fns), with_default_output=with_default_output, with_always_output=with_always_output)

Combine `fns` to execute on _T in parallel across the `fns`.

Args:
        fns: sequence of part functions to chain.
        match_fns: sequence of functions that return True if the part should be
        processed by the part function. When the part should not be processed, the
        part function will not be called and nothing will be yielded by the part
        function. When match_fns is not provided, all parts are processed by
        default.
        with_default_output: True when the parallel execution should fallback to
        return the input part as is when fns do not return any output part.
        with_always_output: True when the parallel execution should always return
        the input part as is independent of the output of the fns. This is a
        stronger condition than `with_default_output`. When `with_always_output`
        is True, `with_default_output` is basically ignored.

Returns:
        Part function that runs all functions 'fns' in parallel. The output stream
        will keep the order of the input parts:

        f_0(c) = c00, c01
        f_1(c) = c10, c11, c12, c14

        f_0(c) // f_1(c) = c00, c01, c10, c11, c12


def resize_attention_map(attentions: tf.Tensor, height: int, width: int, align_corners: bool=False) -> tf.Tensor:
    scale = (height * width // attentions.shape[2]) ** 0.5
    if height > width:
        feat_width = int(np.round(width / scale))
        feat_height = shape_list(attentions)[2] // feat_width
    else:
        feat_height = int(np.round(height / scale))
        feat_width = shape_list(attentions)[2] // feat_height
    batch_size = shape_list(attentions)[0]
    groups = shape_list(attentions)[1]
    attentions = tf.reshape(attentions, (batch_size, groups, feat_height, feat_width))
    attentions = tf.transpose(attentions, perm=(0, 2, 3, 1))
    if align_corners:
        attentions = tf.compat.v1.image.resize(attentions, size=(height, width), method='bilinear', align_corners=align_corners)
    else:
        attentions = tf.image.resize(attentions, size=(height, width), method='bilinear')
    attentions = tf.transpose(attentions, perm=(0, 3, 1, 2))
    return attentions

Args:
    attentions (`tf.Tensor`): attention map of shape [batch_size, groups, feat_height*feat_width]
    height (`int`): height of the output attention map
    width (`int`): width of the output attention map
    align_corners (`bool`, *optional*): the `align_corner` argument for `nn.functional.interpolate`.

Returns:
    `tf.Tensor`: resized attention map of shape [batch_size, groups, height, width]


def waitAndGet(self, event_name, timeout=DEFAULT_TIMEOUT):
    if timeout:
        if timeout > MAX_TIMEOUT:
            raise Error(self._ad, 'Specified timeout %s is longer than max timeout %s.' % (timeout, MAX_TIMEOUT))
    try:
        raw_event = self._callEventWaitAndGet(self._id, event_name, timeout)
    except Exception as e:
        if 'EventSnippetException: timeout.' in str(e):
            raise TimeoutError(self._ad, 'Timed out after waiting %ss for event "%s" triggered by %s (%s).' % (timeout, event_name, self._method_name, self._id))
        raise
    return snippet_event.from_dict(raw_event)

Blocks until an event of the specified name has been received and
    return the event, or timeout.

Args:
        event_name: string, name of the event to get.
        timeout: float, the number of seconds to wait before giving up.

Returns:
        SnippetEvent, the oldest entry of the specified event.

Raises:
        Error: If the specified timeout is longer than the max timeout
        supported.
        TimeoutError: The expected event does not occur within time limit.


def RegisterImplementation(source):
    global _source_implementations
    if 'name' not in source.__dict__:
        raise RuntimeError("'name' not defined in Source %r" % (source,))
    _source_implementations[source.name] = source

Register a Source implementation with the factory method.

    Sources being registered are expected to have a name attribute,
    unique to themselves.

    Child modules are expected to call this method in the file-level
    scope.

Args:
        source: A class type that is a subclass of Source

Returns:
        Nothing

Raises:
        RuntimeError: no 'name' entry in this source.


def infer_graph(inputs: Optional[Set[EventSetNode]], outputs: Set[EventSetNode]) -> Graph:
    graph = Graph()
    graph.outputs.update(outputs)
    pending_nodes: Set[EventSetNode] = outputs.copy()
    done_nodes: Set[EventSetNode] = set()
    missing_nodes: Set[EventSetNode] = set()
    while pending_nodes:
        node = next(iter(pending_nodes))
        pending_nodes.remove(node)
        assert node not in done_nodes
        graph.add_node(node)
        if inputs is not None and node in inputs:
            graph.inputs.add(node)
            continue
        if node.creator is None:
            if inputs is not None:
                missing_nodes.add(node)
            else:
                graph.inputs.add(node)
            continue
        graph.add_operator(node.creator)
        for input_node in node.creator.inputs.values():
            if input_node in done_nodes:
                continue
            pending_nodes.add(input_node)
        for output_node in node.creator.outputs.values():
            graph.add_node(output_node)
    if missing_nodes:
        raise ValueError(f'The following input nodes are required but not provided as input:\n{missing_nodes}')
    for e in graph.nodes:
        graph.add_sampling(e.sampling_node)
        for f in e.feature_nodes:
            graph.add_feature(f)
    return graph

Extracts the nodes in between the output and input nodes.

    If inputs is set, fails if outputs cannot be computed from `inputs`.
    If inputs is not set, infers the required set of inputs.

Args:
        inputs: Set of available input nodes. If None, inputs are inferred.
        outputs: Set of expected output nodes.

Returns:
        The inferred graph.

Raises:
        ValueError: If there are repeated nodes in the `inputs`; an
            unexpected type of input is provided; an unnamed node is inferred
            as input; or some nodes are required but not provided.


def adjoint(self, name: str='adjoint') -> 'LinearOperator':
    if self.is_self_adjoint is True:
        return self
    with self._name_scope(name):
        return self._linop_adjoint()

Returns the adjoint of the current `LinearOperator`.

    Given `A` representing this `LinearOperator`, return `A*`.
    Note that calling `self.adjoint()` and `self.H` are equivalent.

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

Returns:
        `LinearOperator` which represents the adjoint of this `LinearOperator`.


def shape_tensor(self, name='shape_tensor'):
    with self._name_scope(name):
        if self.shape.is_fully_defined():
            return linear_operator_util.shape_tensor(self.shape.as_list())
        else:
            return self._shape_tensor()

Shape of this `LinearOperator`, determined at runtime.

    If this operator acts like the batch matrix `A` with
    `A.shape = [B1,...,Bb, M, N]`, then this returns a `Tensor` holding
    `[B1,...,Bb, M, N]`, equivalent to `tf.shape(A)`.

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

Returns:
        `int32` `Tensor`


def __init__(self, l2=0.01, **kwargs):
    l2 = kwargs.pop('l', l2)
    if kwargs:
        raise TypeError('Argument(s) not recognized: %s' % (kwargs,))
    l2 = 0.01 if l2 is None else l2
    _check_penalty_number(l2)
    self.l2 = backend.cast_to_floatx(l2)

A regularizer that applies a L2 regularization penalty.

    The L2 regularization penalty is computed as:
    `loss = l2 * reduce_sum(square(x))`

    L2 may be passed to a layer as a string identifier:

    >>> dense = tf.keras.layers.Dense(3, kernel_regularizer='l2')

    In this case, the default value used is `l2=0.01`.

Attributes:
        l2: Float; L2 regularization factor.


def GetAutomountMasterMap(self):
    search_base = self.conf['base']
    search_scope = ldap.SCOPE_SUBTREE
    search_filter = '(&(objectclass=automountMap)(ou=auto.master))'
    self.log.debug('retrieving automount master map.')
    self.Search(search_base=search_base, search_filter=search_filter, search_scope=search_scope, attrs=['dn'])
    search_base = None
    for obj in self:
        search_base = obj['dn']
    if search_base is None:
        self.log.critical('Could not find automount master map!')
        raise error.EmptyMap
    self.log.debug('found ou=auto.master at %s', search_base)
    master_map = self.GetAutomountMap(location=search_base)
    for map_entry in master_map:
        map_entry.location = map_entry.location.split(':')[2]
        map_entry.location = map_entry.location.split(' ')[0]
        self.log.debug('master map has: %s' % map_entry.location)
    return master_map

Return the autmount master map from this source.

    The automount master map is a special-case map which points to a dynamic
    list of additional maps. We currently support only the schema outlined at
    http://docs.sun.com/source/806-4251-10/mapping.htm commonly used by linux
    automount clients, namely ou=auto.master and objectclass=automount entries.

Returns:
        an instance of maps.AutomountMap