def node_run(input_file, coords_only, bc_settings, bc_grid_weights):
    log = logging.getLogger('pyspark')
    log.setLevel(logging.INFO)
    if (len(log.handlers) == 0):
        log.addHandler(logging.StreamHandler(sys.stdout))
    precision = bc_settings.value['precision']
    imager = oskar.Imager(precision)
    for (key, value) in bc_settings.value['imager'].items():
        setattr(imager, key, value)
    grid_size = imager.plane_size
    grid_weights = None
    ms_han = oskar.MeasurementSet.open(input_file)
    if coords_only:
        if (imager.weighting == 'Uniform'):
            grid_weights = numpy.zeros([grid_size, grid_size], dtype=precision)
        log.info('Reading coordinates from %s', input_file)
        imager.coords_only = True
        process_input_data(ms_han, imager, None, grid_weights)
        imager.coords_only = False
        return (grid_weights, imager.num_w_planes)
    grid_data = numpy.zeros([grid_size, grid_size], dtype=('c8' if (precision == 'single') else 'c16'))
    log.info('Reading visibilities from %s', input_file)
    if bc_settings.value['combine']:
        if (imager.weighting == 'Uniform'):
            grid_weights = bc_grid_weights.value
        grid_norm = process_input_data(ms_han, imager, grid_data, grid_weights)
        log.info('Returning gridded visibilities to RDD')
        return (grid_data, grid_norm)
    else:
        if (imager.weighting == 'Uniform'):
            grid_weights = numpy.zeros([grid_size, grid_size], dtype=precision)
        if ((imager.weighting == 'Uniform') or (imager.algorithm == 'W-projection')):
            imager.coords_only = True
            process_input_data(ms_han, imager, None, grid_weights)
            imager.coords_only = False
        grid_norm = process_input_data(ms_han, imager, grid_data, grid_weights)
        output_file = (splitext(input_file)[0] + '.fits')
        save_image(imager, grid_data, grid_norm, output_file)
        log.info('Finished. Output file is %s', output_file)
        return 0

Main function to process visibility data on Spark cluster nodes.

Args:
        input_file (str):
            RDD element containing filename to process.
        coords_only (boolean):
            If true, read only baseline coordinates to define the weights grid.
        bc_settings (pyspark.broadcast.Broadcast):
            Spark broadcast variable containing pipeline settings dictionary.
        bc_grid_weights (pyspark.broadcast.Broadcast):
            Spark broadcast variable containing weights grid. May be None.

Returns:
        tuple: Output RDD element.


def sensor(self, name, config=None, inactive_sensor_expiration_time_seconds=sys.maxsize, parents=None):
    sensor = self.get_sensor(name)
    if sensor:
        return sensor
    with self._lock:
        sensor = self.get_sensor(name)
        if (not sensor):
            sensor = Sensor(self, name, parents, (config or self.config), inactive_sensor_expiration_time_seconds)
            self._sensors[name] = sensor
            if parents:
                for parent in parents:
                    children = self._children_sensors.get(parent)
                    if (not children):
                        children = []
                        self._children_sensors[parent] = children
                    children.append(sensor)
            logger.debug('Added sensor with name %s', name)
        return sensor

Get or create a sensor with the given unique name and zero or
    more parent sensors. All parent sensors will receive every value
    recorded with this sensor.

Args:
        name (str): The name of the sensor
        config (MetricConfig, optional): A default configuration to use
            for this sensor for metrics that don't have their own config
        inactive_sensor_expiration_time_seconds (int, optional):
            If no value if recorded on the Sensor for this duration of
            time, it is eligible for removal
        parents (list of Sensor): The parent sensors

Returns:
        Sensor: The sensor that is created


def RemoveScanNode(self, path_spec):
    scan_node = self._scan_nodes.get(path_spec, None)
    if (not scan_node):
        return None
    if scan_node.sub_nodes:
        raise RuntimeError('Scan node has sub nodes.')
    parent_scan_node = scan_node.parent_node
    if parent_scan_node:
        parent_scan_node.sub_nodes.remove(scan_node)
    if (path_spec == self._root_path_spec):
        self._root_path_spec = None
    del self._scan_nodes[path_spec]
    if path_spec.IsFileSystem():
        del self._file_system_scan_nodes[path_spec]
    return parent_scan_node

Removes a scan node of a certain path specification.

Args:
        path_spec (PathSpec): path specification.

Returns:
        SourceScanNode: parent scan node or None if not available.

Raises:
        RuntimeError: if the scan node has sub nodes.


def encode(cls, command):
    args = []
    for arg in command.args:
        if (not isinstance(arg, str)):
            arg = str(arg)
        if ((',' in arg) or arg.startswith(' ') or arg.endswith(' ') or arg.startswith('hex:')):
            arg = 'hex:{}'.format(hexlify(arg.encode('utf-8')).decode('utf-8'))
        args.append(arg)
    argstr = ''
    if (len(args) > 0):
        argstr = ((' {' + ','.join(args)) + '}')
    return (command.name + argstr)

Encode a command as an unambiguous string.

Args:
        command (Command): The command to encode.

Returns:
        str: The encoded command


def snake_to_camel(name):
    ret = ''.join((x.title() for x in name.split('_')))
    ret = (ret[0].lower() + ret[1:])
    return ret

Takes a snake_field_name and returns a camelCaseFieldName

Args:
        name (str): E.g. snake_field_name or SNAKE_FIELD_NAME

Returns:
        str: camelCase converted name. E.g. capsFieldName


def ReadByte(self, do_ord=True):
    try:
        if do_ord:
            return ord(self.stream.read(1))
        return self.stream.read(1)
    except Exception as e:
        logger.error('ord expected character but got none')
    return 0

Read a single byte.

Args:
        do_ord (bool): (default True) convert the byte to an ordinal first.

Returns:
        bytes: a single byte if successful. 0 (int) if an exception occurred.


def make_instance(cls, data):
    schema = cls()
    if (not hasattr(schema.Meta, 'model')):
        raise AttributeError('In order to make an instance, a model for the schema must be defined in the Meta class.')
    serialized_data = schema.load(data).data
    return cls.Meta.model(**serialized_data)

Validate the data and create a model instance from the data.

Args:
        data (dict): The unserialized data to insert into the new model
            instance through it's constructor.

Returns:
        peewee.Model|sqlalchemy.Model: The model instance with it's data
            inserted into it.

Raises:
        AttributeError: This is raised if ``Meta.model`` isn't set on the
            schema's definition.


def pretty_polyfit_plot(x, y, deg=1, xlabel=None, ylabel=None, **kwargs):
    plt = pretty_plot(**kwargs)
    pp = np.polyfit(x, y, deg)
    xp = np.linspace(min(x), max(x), 200)
    plt.plot(xp, np.polyval(pp, xp), 'k--', x, y, 'o')
    if xlabel:
        plt.xlabel(xlabel)
    if ylabel:
        plt.ylabel(ylabel)
    return plt

Convenience method to plot data with trend lines based on polynomial fit.

Args:
        x: Sequence of x data.
        y: Sequence of y data.
        deg (int): Degree of polynomial. Defaults to 1.
        xlabel (str): Label for x-axis.
        ylabel (str): Label for y-axis.
        \\*\\*kwargs: Keyword args passed to pretty_plot.

Returns:
        matplotlib.pyplot object.


def upload_timeline(self, timeline_name, plaso_storage_path):
    resource_url = '{0:s}/upload/'.format(self.api_base_url)
    files = {'file': open(plaso_storage_path, 'rb')}
    data = {'name': timeline_name}
    response = self.session.post(resource_url, files=files, data=data)
    try:
        response_dict = response.json()
    except ValueError:
        raise RuntimeError('Could not decode JSON response from Timesketch (Status {0:d}):\n{1:s}'.format(response.status_code, response.content))
    index_id = response_dict['objects'][0]['id']
    return index_id

Create a timeline with the specified name from the given plaso file.

Args:
        timeline_name (str): Name of timeline
        plaso_storage_path (str): Local path of plaso file to be uploaded

Returns:
        int: ID of uploaded timeline

Raises:
        RuntimeError: When the JSON response from Timesketch cannot be decoded.


def ContainsAddressStr(self, address):
    for (key, contract) in self._contracts.items():
        if (contract.Address == address):
            return True
    return False

Determine if the wallet contains the address.

Args:
        address (str): a string representing the public key.

Returns:
        bool: True, if the address is present in the wallet. False otherwise.


def sheet_to_table(worksheet):
    headers = []
    value_rows = []
    for (row_i, row) in enumerate(worksheet.iter_rows()):
        if (row_i == 0):
            for header_cell in row:
                if header_cell.value:
                    headers.append(parse_value(header_cell))
                else:
                    break
            continue
        row_cells = [parse_value(cell) for (index, cell) in enumerate(row) if (index < len(headers))]
        if any(row_cells):
            value_rows.append(row_cells)
        else:
            break
    table = [{k: v for (k, v) in zip(headers, row) if (v is not None)} for row in value_rows]
    return table

Transforma una hoja de libro de Excel en una lista de diccionarios.

Args:
        worksheet (Workbook.worksheet): Hoja de cálculo de un archivo XLSX
            según los lee `openpyxl`

Returns:
        list_of_dicts: Lista de diccionarios, con tantos elementos como
            registros incluya la hoja, y con tantas claves por diccionario como
            campos tenga la hoja.


def update(self, resource, timeout=(- 1)):
    self.__set_default_values(resource)
    uri = self._client.build_uri(resource['logicalSwitch']['uri'])
    return self._client.update(resource, uri=uri, timeout=timeout)

Updates a Logical Switch.

Args:
        resource (dict): Object to update.
        timeout:
            Timeout in seconds. Wait for task completion by default. The timeout does not abort the operation
            in OneView, just stop waiting for its completion.

Returns:
        dict: Updated resource.


def combine_with_wd_noise(f_n, amp_n, f_n_wd, amp_n_wd):
    amp_n_wd_interp = interpolate.interp1d(f_n_wd, amp_n_wd, bounds_error=False, fill_value=1e-30)
    amp_n_wd = amp_n_wd_interp(f_n)
    amp_n = ((amp_n * (amp_n >= amp_n_wd)) + (amp_n_wd * (amp_n < amp_n_wd)))
    return (f_n, amp_n)

Combine noise with wd noise.

    Combines noise and white dwarf background noise based on greater
    amplitude value at each noise curve step.

Args:
        f_n (float array): Frequencies of noise curve.
        amp_n (float array): Amplitude values of noise curve.
        f_n_wd (float array): Frequencies of wd noise.
        amp_n_wd (float array): Amplitude values of wd noise.

Returns:
        (tuple of float arrays): Amplitude values of combined noise curve.


def _FormatDate(self, event):
    date_time = dfdatetime_posix_time.PosixTimeInMicroseconds(timestamp=event.timestamp)
    (year, month, day_of_month) = date_time.GetDate()
    try:
        return '{0:04d}-{1:02d}-{2:02d}'.format(year, month, day_of_month)
    except (TypeError, ValueError):
        self._ReportEventError(event, 'unable to copy timestamp: {0!s} to a human readable date. Defaulting to: "0000-00-00"'.format(event.timestamp))
        return '0000-00-00'

Formats the date.

Args:
        event (EventObject): event.

Returns:
        str: date field.


def _NormalizeKeyPath(self, key_path):
    normalized_key_path = key_path.lower()
    if ((len(normalized_key_path) < 39) or (not normalized_key_path.startswith(self._CONTROL_SET_PREFIX))):
        return normalized_key_path
    return ''.join([self._NORMALIZED_CONTROL_SET_PREFIX, normalized_key_path[39:]])

Normalizes a Windows Registry key path.

Args:
        key_path (str): Windows Registry key path.

Returns:
        str: normalized Windows Registry key path.


def tags(pode, leaf=False):
    fulltags = [tag for tag in pode[1]['tags']]
    if (not leaf):
        return fulltags
    retn = []
    for (size, tag) in sorted([(len(t), t) for t in fulltags], reverse=True):
        look = (tag + '.')
        if any([r.startswith(look) for r in retn]):
            continue
        retn.append(tag)
    return retn

Get all the tags for a given node.

Args:
        pode (tuple): A packed node.
        leaf (bool): If True, only return the full tags.

Returns:
        list: A list of tag strings.


def dummyctrl(self, r, ctrl):
    dv = DummyVertex(r)
    (dv.view.w, dv.view.h) = (self.dw, self.dh)
    self.grx[dv] = dv
    dv.ctrl = ctrl
    ctrl[r] = dv
    self.layers[r].append(dv)
    return dv

creates a DummyVertex at rank r inserted in the ctrl dict
    of the associated edge and layer.

Args:
        r (int): rank value
        ctrl (dict): the edge's control vertices

Returns:
        DummyVertex : the created DummyVertex.


def moveRel(xOffset=None, yOffset=None, duration=0.0, tween=linear, pause=None, _pause=True):
    _failSafeCheck()
    (xOffset, yOffset) = _unpackXY(xOffset, yOffset)
    _mouseMoveDrag('move', None, None, xOffset, yOffset, duration, tween)
    _autoPause(pause, _pause)

Moves the mouse cursor to a point on the screen, relative to its current
    position.

    The x and y parameters detail where the mouse event happens. If None, the
    current mouse position is used. If a float value, it is rounded down. If
    outside the boundaries of the screen, the event happens at edge of the
    screen.

Args:
        x (int, float, None, tuple, optional): How far left (for negative values) or
        right (for positive values) to move the cursor. 0 by default. If tuple, this is used for x and y.
        y (int, float, None, optional): How far up (for negative values) or
        down (for positive values) to move the cursor. 0 by default.
        duration (float, optional): The amount of time it takes to move the mouse
        cursor to the new xy coordinates. If 0, then the mouse cursor is moved
        instantaneously. 0.0 by default.
        tween (func, optional): The tweening function used if the duration is not
        0. A linear tween is used by default. See the tweens.py file for
        details.

Returns:
        None


def is40(msg):
    if allzeros(msg):
        return False
    d = hex2bin(data(msg))
    if wrongstatus(d, 1, 2, 13):
        return False
    if wrongstatus(d, 14, 15, 26):
        return False
    if wrongstatus(d, 27, 28, 39):
        return False
    if wrongstatus(d, 48, 49, 51):
        return False
    if wrongstatus(d, 54, 55, 56):
        return False
    if (bin2int(d[39:47]) != 0):
        return False
    if (bin2int(d[51:53]) != 0):
        return False
    return True

Check if a message is likely to be BDS code 4,0

Args:
        msg (String): 28 bytes hexadecimal message string

Returns:
        bool: True or False


def map(self, map_fn, desc=None):
    if (desc is None):
        desc = getattr(map_fn, '__name__', '')
    desc = u'map({})'.format(desc)
    return self.transform((lambda xs: (map_fn(x) for x in xs)), desc=desc)

Return a copy of this query, with the values mapped through `map_fn`.

Args:
        map_fn (callable): A callable that takes a single argument and returns a new value.

    Keyword Args:
        desc (str): A description of the mapping transform, for use in log message.
            Defaults to the name of the map function.

Returns:
        Query


def _netsh_file(content):
    with tempfile.NamedTemporaryFile(mode='w', prefix='salt-', suffix='.netsh', delete=False) as fp:
        fp.write(content)
    try:
        log.debug('%s:\n%s', fp.name, content)
        return salt.modules.cmdmod.run('netsh -f {0}'.format(fp.name), python_shell=True)
    finally:
        os.remove(fp.name)

helper function to get the results of ``netsh -f content.txt``

    Running ``netsh`` will drop you into a ``netsh`` prompt where you can issue
    ``netsh`` commands. You can put a series of commands in an external file and
    run them as if from a ``netsh`` prompt using the ``-f`` switch. That's what
    this function does.

Args:
        content (str):
            The contents of the file that will be run by the ``netsh -f``
            command

Returns:
        str: The text returned by the netsh command


def get_leaves(self, item_ids=None, language=None, forbidden_item_ids=None):
    forbidden_item_ids = (set() if (forbidden_item_ids is None) else set(forbidden_item_ids))
    children = self.get_children_graph(item_ids, language=language, forbidden_item_ids=forbidden_item_ids)
    counts = self.get_children_counts(active=None)
    if (item_ids is None):
        item_ids = set(children.keys())

    def _get_leaves(item_id):
        leaves = set()

        def __search(item_ids):
            result = set(flatten([children.get(item_id, []) for item_id in item_ids]))
            new_leaves = {item_id for item_id in result if (item_id not in children.keys())}
            leaves.update(new_leaves)
            return (result - new_leaves)
        fixed_point(is_zero=(lambda to_visit: (len(to_visit) == 0)), minus=(lambda to_visit, visited: (to_visit - visited)), plus=(lambda visited_x, visited_y: (visited_x | visited_y)), f=__search, x={item_id})
        leaves = {leaf for leaf in leaves if (counts[leaf] == 0)}
        if (len(leaves) > 0):
            return leaves
        if ((counts[item_id] == 0) and (item_id not in forbidden_item_ids)):
            return {item_id}
        return set()
    return {item_id: _get_leaves(item_id) for item_id in item_ids}

Get mapping of items to their reachable leaves. Leaves having
    inactive relations to other items are omitted.

Args:
        item_ids (list): items which are taken as roots for the reachability
        language (str): if specified, filter out items which are not
            available in the given language

Returns:
        dict: item id -> list of items (reachable leaves)


def _get_endpoint(self, sub_domain):
    storage_parameters = (self._storage_parameters or dict())
    account_name = storage_parameters.get('account_name')
    if (not account_name):
        raise ValueError('"account_name" is required for Azure storage')
    suffix = storage_parameters.get('endpoint_suffix', 'core.windows.net')
    self._endpoint = ('http%s://%s.%s.%s' % (('' if self._unsecure else 's'), account_name, sub_domain, suffix))
    return (account_name, suffix.replace('.', '\\.'))

Get endpoint information from storage parameters.

    Update system with endpoint information and return information required
    to define roots.

Args:
        self (pycosio._core.io_system.SystemBase subclass): System.
        sub_domain (str): Azure storage sub-domain.

Returns:
        tuple of str: account_name, endpoint_suffix


def ReadLine(self, file_object):
    (line, _, self.lines) = self.lines.partition('\n')
    if (not line):
        self.ReadLines(file_object)
        (line, _, self.lines) = self.lines.partition('\n')
    return line

Reads a line.

Args:
        file_object (dfvfs.FileIO): file-like object.

Returns:
        str: line read from the lines buffer.


def build_hgnc_gene(gene_info, build='37'):
    try:
        hgnc_id = int(gene_info['hgnc_id'])
    except KeyError as err:
        raise KeyError('Gene has to have a hgnc_id')
    except ValueError as err:
        raise ValueError('hgnc_id has to be integer')
    try:
        hgnc_symbol = gene_info['hgnc_symbol']
    except KeyError as err:
        raise KeyError('Gene has to have a hgnc_symbol')
    try:
        ensembl_id = gene_info['ensembl_gene_id']
    except KeyError as err:
        raise KeyError('Gene has to have a ensembl_id')
    try:
        chromosome = gene_info['chromosome']
    except KeyError as err:
        raise KeyError('Gene has to have a chromosome')
    try:
        start = int(gene_info['start'])
    except KeyError as err:
        raise KeyError('Gene has to have a start position')
    except TypeError as err:
        raise TypeError('Gene start has to be a integer')
    try:
        end = int(gene_info['end'])
    except KeyError as err:
        raise KeyError('Gene has to have a end position')
    except TypeError as err:
        raise TypeError('Gene end has to be a integer')
    gene_obj = HgncGene(hgnc_id=hgnc_id, hgnc_symbol=hgnc_symbol, ensembl_id=ensembl_id, chrom=chromosome, start=start, end=end, build=build)
    if gene_info.get('description'):
        gene_obj['description'] = gene_info['description']
    if gene_info.get('previous_symbols'):
        gene_obj['aliases'] = gene_info['previous_symbols']
    if gene_info.get('entrez_id'):
        gene_obj['entrez_id'] = int(gene_info['entrez_id'])
    if gene_info.get('omim_id'):
        gene_obj['omim_id'] = int(gene_info['omim_id'])
    if gene_info.get('pli_score'):
        gene_obj['pli_score'] = float(gene_info['pli_score'])
    if gene_info.get('ref_seq'):
        gene_obj['primary_transcripts'] = gene_info['ref_seq']
    if gene_info.get('ucsc_id'):
        gene_obj['ucsc_id'] = gene_info['ucsc_id']
    if gene_info.get('uniprot_ids'):
        gene_obj['uniprot_ids'] = gene_info['uniprot_ids']
    if gene_info.get('vega_id'):
        gene_obj['vega_id'] = gene_info['vega_id']
    if gene_info.get('incomplete_penetrance'):
        gene_obj['incomplete_penetrance'] = True
    if gene_info.get('inheritance_models'):
        gene_obj['inheritance_models'] = gene_info['inheritance_models']
    phenotype_objs = []
    for phenotype_info in gene_info.get('phenotypes', []):
        phenotype_objs.append(build_phenotype(phenotype_info))
    if phenotype_objs:
        gene_obj['phenotypes'] = phenotype_objs
    for key in list(gene_obj):
        if (gene_obj[key] is None):
            gene_obj.pop(key)
    return gene_obj

Build a hgnc_gene object

Args:
        gene_info(dict): Gene information

Returns:
        gene_obj(dict)

        {
            '_id': ObjectId(),
            # This is the hgnc id, required:
            'hgnc_id': int,
            # The primary symbol, required
            'hgnc_symbol': str,
            'ensembl_id': str, # required
            'build': str, # '37' or '38', defaults to '37', required

            'chromosome': str, # required
            'start': int, # required
            'end': int, # required

            'description': str, # Gene description
            'aliases': list(), # Gene symbol aliases, includes hgnc_symbol, str
            'entrez_id': int,
            'omim_id': int,
            'pli_score': float,
            'primary_transcripts': list(), # List of refseq transcripts (str)
            'ucsc_id': str,
            'uniprot_ids': list(), # List of str
            'vega_id': str,
            'transcripts': list(), # List of hgnc_transcript

            # Inheritance information
            'inheritance_models': list(), # List of model names
            'incomplete_penetrance': bool, # Acquired from HPO

            # Phenotype information
            'phenotypes': list(), # List of dictionaries with phenotype information
        }


def moveTo(x=None, y=None, duration=0.0, tween=linear, pause=None, _pause=True):
    _failSafeCheck()
    (x, y) = _unpackXY(x, y)
    _mouseMoveDrag('move', x, y, 0, 0, duration, tween)
    _autoPause(pause, _pause)

Moves the mouse cursor to a point on the screen.

    The x and y parameters detail where the mouse event happens. If None, the
    current mouse position is used. If a float value, it is rounded down. If
    outside the boundaries of the screen, the event happens at edge of the
    screen.

Args:
        x (int, float, None, tuple, optional): The x position on the screen where the
        click happens. None by default. If tuple, this is used for x and y.
        If x is a str, it's considered a filename of an image to find on
        the screen with locateOnScreen() and click the center of.
        y (int, float, None, optional): The y position on the screen where the
        click happens. None by default.
        duration (float, optional): The amount of time it takes to move the mouse
        cursor to the xy coordinates. If 0, then the mouse cursor is moved
        instantaneously. 0.0 by default.
        tween (func, optional): The tweening function used if the duration is not
        0. A linear tween is used by default. See the tweens.py file for
        details.

Returns:
        None


def wait_for_vacancy(self, processor_type):
    with self._condition:
        self._condition.wait_for((lambda : (self._processor_available(processor_type) or self._cancelled_event.is_set())))
        if self._cancelled_event.is_set():
            raise WaitCancelledException()
        processor = self[processor_type].next_processor()
        return processor

Waits for a particular processor type to have the capacity to
    handle additional transactions or until is_cancelled is True.

Args:
        processor_type (ProcessorType): The family, and version of
            the transaction processor.

Returns:
        Processor


def broadcast_dynamic_shape(shape_x: DynamicRaggedShape, shape_y: DynamicRaggedShape) -> DynamicRaggedShape:
    if not isinstance(shape_x, DynamicRaggedShape):
        raise TypeError('shape_x must be a DynamicRaggedShape')
    if not isinstance(shape_y, DynamicRaggedShape):
        raise TypeError('shape_y must be a DynamicRaggedShape')
    return broadcast_dynamic_shape_extended(shape_x, shape_y)[0]

Returns the shape formed by broadcasting two shapes to be compatible.

    1. If shape_x and shape_y both have row_partitions, then fail if their dtypes
        don't match.
    2. If neither has row_partitions and they have different dtypes,
        go with int64.
    3. If one has row_partitions, go with that dtype.

Args:
        shape_x: A `DynamicRaggedShape`
        shape_y: A `DynamicRaggedShape`

Returns:
        A `DynamicRaggedShape`.

Raises:
        ValueError: If `shape_x` and `shape_y` are not broadcast-compatible.


def register_revived_type(identifier, predicate, versions):
    versions.sort(key=lambda reg: reg.version, reverse=True)
    if not versions:
        raise AssertionError('Need at least one version of a registered type.')
    version_numbers = set()
    for registration in versions:
        registration.identifier = identifier
        if registration.version in version_numbers:
            raise AssertionError(f'Got multiple registrations with version {registration.version} for type {identifier}.')
        version_numbers.add(registration.version)
    _REVIVED_TYPE_REGISTRY[identifier] = (predicate, versions)
    _TYPE_IDENTIFIERS.append(identifier)

Register a type for revived objects.

Args:
        identifier: A unique string identifying this class of objects.
        predicate: A Boolean predicate for this registration. Takes a
        trackable object as an argument. If True, `type_registration` may be
        used to save and restore the object.
        versions: A list of `VersionedTypeRegistration` objects.


def check_subword_sampling(tokenizer: PreTrainedTokenizer, text: Optional[str]=None, test_sentencepiece_ignore_case: bool=True) -> None:
    text = 'This is a test for subword regularization.' if text is None else text
    if test_sentencepiece_ignore_case:
        text = text.lower()
    tokens_list = []
    for _ in range(5):
        tokens_list.append(tokenizer.tokenize(text))
    combinations = itertools.combinations(tokens_list, 2)
    subword_sampling_found = False
    for combination in combinations:
        if combination[0] != combination[1]:
            subword_sampling_found = True
    unittest.TestCase().assertTrue(subword_sampling_found)
    for tokens in tokens_list:
        if test_sentencepiece_ignore_case:
            unittest.TestCase().assertEqual(text, tokenizer.convert_tokens_to_string(tokens).lower())
        else:
            unittest.TestCase().assertEqual(text, tokenizer.convert_tokens_to_string(tokens))

Check if the tokenizer generates different results when subword regularization is enabled.

    Subword regularization augments training data with subword sampling.
    This has a random component.

Args:
        tokenizer: The tokenizer to check.
        text: The text to use for the checks.
        test_sentencepiece_ignore_case: See `TokenizerTesterMixin.test_sentencepiece_ignore_case`.


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

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

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

    Subclasses should override for any actions to run.

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

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


def __init__(self, query_engine: engine.Engine, fhir_dataset: str, view_dataset: Optional[str]=None, value_set_codes_table: Optional[str]=None, snake_case_resource_tables: bool=False) -> None:
    self._engine = query_engine
    self._fhir_dataset = fhir_dataset
    self._view_dataset = view_dataset if view_dataset is not None else self._fhir_dataset
    self._value_set_codes_table = value_set_codes_table if value_set_codes_table is not None else 'value_set_codes'
    self._snake_case_resource_tables = snake_case_resource_tables
    self._value_set_manager = spark_value_set_manager.SparkValueSetManager(query_engine, self._value_set_codes_table)

Initializes the SparkRunner with provided SQLAlcehmy Engine and Dataset.

Args:
        query_engine: SQLAlchemy Engine with which to perform queries.
        fhir_dataset: Dataset with FHIR data that the views will query.
        view_dataset: Optional dataset with views will be created via the
        `to_spark_view` method, if used. It will use the fhir_dataset if no
        view_dataset is specified
        value_set_codes_table: A table containing value set expansions. If
        provided, memberOf queries may be made against value set URLs described
        by this table. If `value_set_codes_table` is a string, it must included
        a project ID if not in the client's default project, dataset ID and
        table ID, each separated by a '.'. The table must match the schema
        described by
        https://github.com/FHIR/sql-on-fhir/blob/master/sql-on-fhir.md#valueset-support
        snake_case_resource_tables: Whether to use snake_case names for resource
        tables in Spark for compatiblity with some exports. Defaults to
        False.


def __init__(self, partitioned_dim_sizes, inner_dim_sizes, dim_size_dtype=None):
    assert isinstance(partitioned_dim_sizes, (list, tuple))
    with ops.name_scope(None, 'RaggedTensorDynamicShape', (partitioned_dim_sizes, inner_dim_sizes)):
        partitioned_dim_sizes = tuple((ops.convert_to_tensor(size, name='partitioned_dimension_size_%d' % i) for i, size in enumerate(partitioned_dim_sizes)))
        inner_dim_sizes = ops.convert_to_tensor(inner_dim_sizes, name='inner_dim_sizes')
        if partitioned_dim_sizes:
            for axis, dimension_size in enumerate(partitioned_dim_sizes):
                if dimension_size.shape.ndims is None:
                    raise ValueError('rank of partitioned_dim_sizes[%d] is unknown' % axis)
                dimension_size.shape.with_rank_at_most(1)
            if partitioned_dim_sizes[0].shape.ndims == 1:
                raise ValueError('outermost partitioned dimension must be uniform')
            if partitioned_dim_sizes[-1].shape.ndims == 0:
                raise ValueError('innermost partitioned dimension must be ragged')
        inner_dim_sizes.shape.assert_has_rank(1)
        if dim_size_dtype is None:
            dim_size_dtypes = set((p.dtype for p in partitioned_dim_sizes if p.shape.ndims == 1))
            if not dim_size_dtypes:
                dim_size_dtype = dtypes.int64
            elif len(dim_size_dtypes) == 1:
                dim_size_dtype = dim_size_dtypes.pop()
            else:
                if not ragged_config.auto_cast_partition_dtype():
                    raise ValueError('partitioned_dim_sizes must have matching dtypes')
                dim_size_dtype = dtypes.int64
        partitioned_dim_sizes = tuple((math_ops.cast(p, dim_size_dtype) for p in partitioned_dim_sizes))
        inner_dim_sizes = math_ops.cast(inner_dim_sizes, dim_size_dtype)
        self._partitioned_dim_sizes = partitioned_dim_sizes
        self._inner_dim_sizes = inner_dim_sizes

Creates a RaggedTensorDynamicShape.

Args:
        partitioned_dim_sizes: A `list` of 0-D or 1-D integer `Tensor`, one for
        each partitioned dimension.  If dimension `d` is uniform, then
        `partitioned_dim_sizes[d]` must be an integer scalar, specifying the
        size of all slices across dimension `d`.  If dimension `d` is ragged,
        then `partitioned_dim_sizes[d]` must be an integer vector, specifying
        the size of each slice across dimension `d`.
        inner_dim_sizes: A 1-D integer `Tensor`, whose length is equal to the
        number of inner dimensions.  `inner_dim_sizes[n]` is the size of all
        slices across the `n`th inner dimension (which is the
        `(len(partitioned_dim_sizes)+n)`th dimension in the overall tensor.
        dim_size_dtype: dtype for dimension sizes.  If not specified, then it
        is chosen based on the dtypes of `partitioned_dim_sizes` and
        `inner_dim_sizes`.


def __init__(self, alpha=1.0, **kwargs):
    super().__init__(**kwargs)
    self.alpha = alpha
    self.supports_masking = True
    self._build_at_init()

Applies an Exponential Linear Unit function to an output.

    Formula:

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

Args:
        alpha: float, slope of negative section. Defaults to `1.0`.
        **kwargs: Base layer keyword arguments, such as `name` and `dtype`.


def _generate_graph_update_dicts(self):
    transform_dict = {}
    pcoll_dict = {}
    for transform_id, transform_proto in self._top_level_transforms():
        transform_dict[transform_proto.unique_name] = {'required': transform_id in self._required_transforms}
        for pcoll_id in transform_proto.outputs.values():
            pcoll_dict[pcoll_id] = {'cached': pcoll_id in self._cached_pcollections, 'referenced': pcoll_id in self._referenced_pcollections}

    def vertex_properties_to_attributes(vertex):
        """Converts PCollection properties to DOT vertex attributes."""
        attrs = {}
        if 'leaf' in vertex:
            attrs['style'] = 'invis'
        elif vertex.get('required'):
            attrs['color'] = 'blue'
            attrs['fontcolor'] = 'blue'
        else:
            attrs['color'] = 'grey'
        return attrs

    def edge_properties_to_attributes(edge):
        """Converts PTransform properties to DOT edge attributes."""
        attrs = {}
        if edge.get('cached'):
            attrs['color'] = 'red'
        elif edge.get('referenced'):
            attrs['color'] = 'black'
        else:
            attrs['color'] = 'grey'
        return attrs
    vertex_dict = {}
    edge_dict = {}
    for transform_name, transform_properties in transform_dict.items():
        vertex_dict[transform_name] = vertex_properties_to_attributes(transform_properties)
    for pcoll_id, pcoll_properties in pcoll_dict.items():
        edge_dict[pcoll_id] = edge_properties_to_attributes(pcoll_properties)
    return (vertex_dict, edge_dict)

Generate updates specific to interactive pipeline.

Returns:
        vertex_dict: (Dict[str, Dict[str, str]]) maps vertex name to attributes
        edge_dict: (Dict[str, Dict[str, str]]) maps vertex name to attributes


def _populate_quantization_options_default_values(quantization_options: _QuantizationOptions) -> None:
    if quantization_options.op_set == quant_opts_pb2.OpSet.OP_SET_UNSPECIFIED:
        quantization_options.op_set = quant_opts_pb2.OpSet.XLA
    if not quantization_options.tags:
        quantization_options.tags.append(tag_constants.SERVING)
    if not quantization_options.signature_keys:
        quantization_options.signature_keys.append(signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY)
    if not quantization_options.HasField('freeze_all_variables'):
        quantization_options.freeze_all_variables = True
    if quantization_options.enable_legacy_weight_only:
        raise ValueError('Legacy weight-only is deprecated. Use weight-only quantization method.')
    if quantization_options.quantization_method.preset_method == _PresetMethod.METHOD_UNSPECIFIED:
        logging.debug('"preset_method" for QuantizationMethod is not specified.Static range quantization is used by default.')
        quantization_options.quantization_method.preset_method = _PresetMethod.METHOD_STATIC_RANGE_INT8
    if quantization_options.min_num_elements_for_weights == 0:
        quantization_options.min_num_elements_for_weights = _DYNAMIC_RANGE_DEFAULT_MIN_NUM_ELEMENTS_FOR_WEIGHTS
        logging.warning('QuantizationOptions.min_num_elements_for_weights is not set (0). Setting to the default value: %d.', _DYNAMIC_RANGE_DEFAULT_MIN_NUM_ELEMENTS_FOR_WEIGHTS)
    if not quantization_options.HasField('enable_per_channel_quantization'):
        quantization_options.enable_per_channel_quantization = False
    if quantization_options.enable_per_channel_quantization and (not ((quantization_options.op_set == quant_opts_pb2.OpSet.UNIFORM_QUANTIZED or quantization_options.quantization_method.preset_method == _PresetMethod.METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8) or (quantization_options.op_set in (quant_opts_pb2.OpSet.XLA, quant_opts_pb2.OpSet.STABLEHLO) and quantization_options.quantization_method.preset_method == _PresetMethod.METHOD_STATIC_RANGE_INT8))):
        raise ValueError('Currently, per-channel quantization is supported for Uniform Quantized opset, weight only quantization, or XLA/StableHLO opset with static range quantization.')
    if quantization_options.quantization_method.preset_method == _PresetMethod.METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8 and (quantization_options.op_set == quant_opts_pb2.OpSet.UNIFORM_QUANTIZED or quantization_options.op_set == quant_opts_pb2.OpSet.TF):
        raise ValueError('TF/Uniform quantized opset does not support weight-only.')
    if quantization_options.op_set == quant_opts_pb2.OpSet.STABLEHLO and (quantization_options.quantization_method.preset_method != _PresetMethod.METHOD_STATIC_RANGE_INT8 and quantization_options.quantization_method.preset_method != _PresetMethod.METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8):
        raise ValueError('StableHLO quantized opset currently only supports static range quantization and weight-only quantizationvia TF Quantizer.')
    logging.debug('Setting `force_graph_mode_calibration = True` to ensure the calibration mode is executed properly.')
    quantization_options.force_graph_mode_calibration = True
    if quantization_options.HasField('debugger_config'):
        if not quantization_options.debugger_config.log_dir_path:
            quantization_options.debugger_config.log_dir_path = '/tmp/dumps'
        if quantization_options.debugger_config.debugger_type == stablehlo_quant_config_pb2.DebuggerConfig.DebuggerType.DEBUGGER_TYPE_UNSPECIFIED:
            raise ValueError('Debugger is enabled but debugger type was not specified.')
        if quantization_options.debugger_config.debugger_type == stablehlo_quant_config_pb2.DebuggerConfig.DebuggerType.DEBUGGER_TYPE_WHOLE_MODEL and (not quantization_options.debugger_config.unquantized_dump_model_path):
            raise ValueError('Debugger type whole model verify was used but unquantized_dump_model_path was not specified.')
    _populate_quantization_component_spec(quantization_options.quantization_method)
    _populate_unitwise_quantization_specs(quantization_options)
    if quantization_options.quantization_method.preset_method == _PresetMethod.METHOD_STATIC_RANGE_INT8:
        _populate_calibration_options(quantization_options)

Populates default values for QuantizationOptions.

    Populates unspecified or unset fields of QuantizationOptions with the default
    values.

    * If `op_set` is unspecified, it defaults to `OpSet.XLA`.
    * If `freeze_all_variables` is not set, it defaults to `True`.
    * Check if configurations are set correctly:
        - Per-channel quantization is supported for Uniform Quantized opset only.

Args:
        quantization_options: An instance of QuantizationOptions.


def make_input_fn_iterator(self, input_fn, replication_mode=InputReplicationMode.PER_WORKER):
    return super(StrategyV1, self).make_input_fn_iterator(input_fn, replication_mode)

Returns an iterator split across replicas created from an input function.

    DEPRECATED: This method is not available in TF 2.x.

    The `input_fn` should take an `tf.distribute.InputContext` object where
    information about batching and input sharding can be accessed:

    ```
    def input_fn(input_context):
        batch_size = input_context.get_per_replica_batch_size(global_batch_size)
        d = tf.data.Dataset.from_tensors([[1.]]).repeat().batch(batch_size)
        return d.shard(input_context.num_input_pipelines,
                    input_context.input_pipeline_id)
    with strategy.scope():
        iterator = strategy.make_input_fn_iterator(input_fn)
        replica_results = strategy.experimental_run(replica_fn, iterator)
    ```

    The `tf.data.Dataset` returned by `input_fn` should have a per-replica
    batch size, which may be computed using
    `input_context.get_per_replica_batch_size`.

Args:
        input_fn: A function taking a `tf.distribute.InputContext` object and
        returning a `tf.data.Dataset`.
        replication_mode: an enum value of `tf.distribute.InputReplicationMode`.
        Only `PER_WORKER` is supported currently, which means there will be
        a single call to `input_fn` per worker. Replicas will dequeue from the
        local `tf.data.Dataset` on their worker.

Returns:
        An iterator object that should first be `.initialize()`-ed. It may then
        either be passed to `strategy.experimental_run()` or you can
        `iterator.get_next()` to get the next value to pass to
        `strategy.extended.call_for_each_replica()`.


def scatter_min(self, sparse_delta, use_locking=False, name=None):
    raise NotImplementedError

Updates this variable with the min of `tf.IndexedSlices` and itself.

Args:
        sparse_delta: `tf.IndexedSlices` to use as an argument of min with this
        variable.
        use_locking: If `True`, use locking during the operation.
        name: the name of the operation.

Returns:
        The updated variable.

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


def convert_sparse_cross_attention_mask_to_dense(cross_attention_token_mask: List[List[List[int]]], num_tiles: List[List[int]], max_num_tiles: int, length: int) -> np.ndarray:
    batch_size = len(cross_attention_token_mask)
    max_num_images = max([len(masks) for masks in cross_attention_token_mask])
    cross_attention_mask = np.zeros(shape=(batch_size, length, max_num_images, max_num_tiles), dtype=np.int64)
    for sample_idx, (sample_masks, sample_num_tiles) in enumerate(zip(cross_attention_token_mask, num_tiles)):
        for mask_idx, (locations, mask_num_tiles) in enumerate(zip(sample_masks, sample_num_tiles)):
            if len(locations) == 2:
                start, end = locations
                end = min(end, length)
                if end == -1:
                    end = length
                cross_attention_mask[sample_idx, start:end, mask_idx, :mask_num_tiles] = 1
    return cross_attention_mask

Convert the cross attention mask indices to a cross attention mask 4D array.

    This function takes a sparse representation of cross attention masks and converts it to a dense 4D numpy array.
    The sparse representation is a nested list structure that defines attention ranges for each image in each batch item.

Args:
        cross_attention_token_mask (List[List[List[int]]]): A nested list structure where:
            - The outer list represents the batch dimension.
            - The middle list represents different images within each batch item.
            - The inner list contains pairs of integers [start, end] representing token ranges for each image.
        num_tiles (List[List[int]]): A nested list structure specifying the number of tiles for each image in each batch item.
        max_num_tiles (int): The maximum possible number of tiles.
        length (int): The total sequence length of the input.

Returns:
        np.ndarray: A 4D numpy array of shape (batch_size, length, max_num_images, max_num_tiles)
            The array contains `1` where attention is allowed and `0` where it is not.

Note:
        - Special handling is done for cases where the end token is -1, which is interpreted as attending to the end of the sequence.


def index_subdirectory(directory, class_indices, follow_links, formats):
    dirname = os.path.basename(directory)
    valid_files = iter_valid_files(directory, follow_links, formats)
    labels = []
    filenames = []
    for root, fname in valid_files:
        labels.append(class_indices[dirname])
        absolute_path = tf.io.gfile.join(root, fname)
        relative_path = tf.io.gfile.join(dirname, os.path.relpath(absolute_path, directory))
        filenames.append(relative_path)
    return (filenames, labels)

Recursively walks directory and list image paths and their class index.

Args:
        directory: string, target directory.
        class_indices: dict mapping class names to their index.
        follow_links: boolean, whether to recursively follow subdirectories
            (if False, we only list top-level images in `directory`).
        formats: Allowlist of file extensions to index (e.g. ".jpg", ".txt").

Returns:
        tuple `(filenames, labels)`. `filenames` is a list of relative file
            paths, and `labels` is a list of integer labels corresponding
            to these files.


def stack(list_or_tensor, element_dtype=None, strict=True):
    if strict:

        def raise_error(x):
            raise ValueError('%s must be stackable when strict=True' % x)
        original_call = raise_error
    else:
        original_call = lambda x: x
    return data_structures.list_stack(list_or_tensor, data_structures.ListStackOpts(element_dtype=element_dtype, original_call=original_call))

Stacks the input, if it admits the notion of stacking.

    For example, a list of tensors can be stacked into a larger tensor. This
    function is similar to tf.stack, but it accepts non-lists and lists of
    non-tensors as arguments. In the latter case, the function does nothing.

Args:
        list_or_tensor: Any
        element_dtype: tf.DType, optional dtypedtype for the elements in the list.
        Required if the input is stackable, and the list is untyped.
        strict: bool, if True an error is raised if the input is not stackable.
        Otherwise the function is a no-op.

Returns:
        Any, if the input is stackable, the result will be a tf.Tensor. Otherwise,
        if strict=False, the result will be list_or_tensor.

Raises:
        ValueError: if strict=True and the input is not stackable.


def __init__(self, operators, is_non_singular=None, is_self_adjoint=None, is_positive_definite=None, is_square=None, name=None):
    parameters = dict(operators=operators, is_non_singular=is_non_singular, is_self_adjoint=is_self_adjoint, is_positive_definite=is_positive_definite, is_square=is_square, name=name)
    check_ops.assert_proper_iterable(operators)
    operators = list(operators)
    if not operators:
        raise ValueError('Expected a non-empty list of operators. Found: %s' % operators)
    self._operators = operators
    dtype = operators[0].dtype
    for operator in operators:
        if operator.dtype != dtype:
            name_type = (str((o.name, o.dtype)) for o in operators)
            raise TypeError('Expected all operators to have the same dtype.  Found %s' % '   '.join(name_type))
    if all((operator.is_non_singular for operator in operators)):
        if is_non_singular is False:
            raise ValueError('The composition of non-singular operators is always non-singular.')
        is_non_singular = True
    if _composition_must_be_self_adjoint(operators):
        if is_self_adjoint is False:
            raise ValueError('The composition was determined to be self-adjoint but user provided incorrect `False` hint.')
        is_self_adjoint = True
    if linear_operator_util.is_aat_form(operators):
        if is_square is False:
            raise ValueError('The composition was determined have the form A @ A.H, hence it must be square. The user provided an incorrect `False` hint.')
        is_square = True
    if linear_operator_util.is_aat_form(operators) and is_non_singular:
        if is_positive_definite is False:
            raise ValueError('The composition was determined to be non-singular and have the form A @ A.H, hence it must be positive-definite. The user provided an incorrect `False` hint.')
        is_positive_definite = True
    if name is None:
        name = '_o_'.join((operator.name for operator in operators))
    with ops.name_scope(name):
        super(LinearOperatorComposition, self).__init__(dtype=dtype, is_non_singular=is_non_singular, is_self_adjoint=is_self_adjoint, is_positive_definite=is_positive_definite, is_square=is_square, parameters=parameters, name=name)

Initialize a `LinearOperatorComposition`.

    `LinearOperatorComposition` is initialized with a list of operators
    `[op_1,...,op_J]`.  For the `matmul` method to be well defined, the
    composition `op_i.matmul(op_{i+1}(x))` must be defined.  Other methods have
    similar constraints.

Args:
        operators:  Iterable of `LinearOperator` objects, each with
        the same `dtype` and composable shape.
        is_non_singular:  Expect that this operator is non-singular.
        is_self_adjoint:  Expect that this operator is equal to its hermitian
        transpose.
        is_positive_definite:  Expect that this operator is positive definite,
        meaning the quadratic form `x^H A x` has positive real part for all
        nonzero `x`.  Note that we do not require the operator to be
        self-adjoint to be positive-definite.  See:
        https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
        is_square:  Expect that this operator acts like square [batch] matrices.
        name: A name for this `LinearOperator`.  Default is the individual
        operators names joined with `_o_`.

Raises:
        TypeError:  If all operators do not have the same `dtype`.
        ValueError:  If `operators` is empty.


def call(self, input_ids: tf.Tensor | None=None, bbox: tf.Tensor | None=None, attention_mask: tf.Tensor | None=None, token_type_ids: tf.Tensor | None=None, position_ids: tf.Tensor | None=None, head_mask: tf.Tensor | None=None, inputs_embeds: tf.Tensor | None=None, pixel_values: tf.Tensor | None=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, training: bool=False) -> Union[TFBaseModelOutput, Tuple[tf.Tensor], Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Tensor, tf.Tensor, tf.Tensor]]:
    outputs = self.layoutlmv3(input_ids=input_ids, bbox=bbox, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, pixel_values=pixel_values, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, training=training)
    return outputs

Returns:

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

    >>> processor = AutoProcessor.from_pretrained("microsoft/layoutlmv3-base", apply_ocr=False)
    >>> model = TFAutoModel.from_pretrained("microsoft/layoutlmv3-base")

    >>> dataset = load_dataset("nielsr/funsd-layoutlmv3", split="train", trust_remote_code=True)
    >>> example = dataset[0]
    >>> image = example["image"]
    >>> words = example["tokens"]
    >>> boxes = example["bboxes"]

    >>> encoding = processor(image, words, boxes=boxes, return_tensors="tf")

    >>> outputs = model(**encoding)
    >>> last_hidden_states = outputs.last_hidden_state
    ```


def _init_summary_op(self, summary_op=USE_DEFAULT):
    if summary_op is Supervisor.USE_DEFAULT:
        summary_op = self._get_first_op_from_collection(ops.GraphKeys.SUMMARY_OP)
        if summary_op is None:
            summary_op = _summary.merge_all()
            if summary_op is not None:
                ops.add_to_collection(ops.GraphKeys.SUMMARY_OP, summary_op)
    self._summary_op = summary_op

Initializes summary_op.

Args:
        summary_op: An Operation that returns a Summary for the event logs. If set
        to USE_DEFAULT, create an op that merges all the summaries.


def forward(self, inputs_embeds=None, attention_mask=None, object_queries=None, output_attentions=None, output_hidden_states=None, return_dict=None):
    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
    hidden_states = inputs_embeds
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    if attention_mask is not None:
        attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype)
    encoder_states = () if output_hidden_states else None
    all_attentions = () if output_attentions else None
    for encoder_layer in self.layers:
        if output_hidden_states:
            encoder_states = encoder_states + (hidden_states,)
        to_drop = False
        if self.training:
            dropout_probability = torch.rand([])
            if dropout_probability < self.layerdrop:
                to_drop = True
        if to_drop:
            layer_outputs = (None, None)
        else:
            layer_outputs = encoder_layer(hidden_states, attention_mask, object_queries=object_queries, output_attentions=output_attentions)
            hidden_states = layer_outputs[0]
        if output_attentions:
            all_attentions = all_attentions + (layer_outputs[1],)
    if output_hidden_states:
        encoder_states = encoder_states + (hidden_states,)
    hidden_states = self.layernorm(hidden_states)
    if not return_dict:
        return tuple((v for v in [hidden_states, encoder_states, all_attentions] if v is not None))
    return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions)

Args:
    inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
        Flattened feature map (output of the backbone + projection layer) that is passed to the encoder.

    attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
        Mask to avoid performing attention on padding pixel features. Mask values selected in `[0, 1]`:

        - 1 for pixel features that are real (i.e. **not masked**),
        - 0 for pixel features that are padding (i.e. **masked**).

        [What are attention masks?](../glossary#attention-mask)

    object_queries (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
        Position embeddings that are added to the queries and keys in each self-attention layer.

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


def conv_input_length(output_length, filter_size, padding, stride):
    if output_length is None:
        return None
    assert padding in {'same', 'valid', 'full'}
    if padding == 'same':
        pad = filter_size // 2
    elif padding == 'valid':
        pad = 0
    elif padding == 'full':
        pad = filter_size - 1
    return (output_length - 1) * stride - 2 * pad + filter_size

Determines input length of a convolution given output length.

Args:
        output_length: integer.
        filter_size: integer.
        padding: one of "same", "valid", "full".
        stride: integer.

Returns:
        The input length (integer).


def expectation(self, function):
    return self._expectation(function)

Returns an estimate of the expectation value of the given function.

Args:
        function: Mapping from a 2D tensor of bitstrings to a possibly nested
        structure.  The structure must have atomic elements all of which are
        float tensors with the same batch size as the input bitstrings.


def get_device_locations(mesh: layout_lib.Mesh, client_id: Optional[int]=None) -> List[Dict[str, int]]:
    if mesh.device_type() != _TPU_DEVICE_TYPE:
        raise ValueError('The mesh must be a TPU mesh')
    if client_id is None or client_id == config.client_id():
        return mesh.local_device_locations()
    raise NotImplementedError("Looking up other clients' device locations is not supported")

Returns the device locations of all TPU cores local to the given client.

    A device location is a dictionary from dimension names to indices on those
    dimensions. For example, for a 2x2 mesh ('x', 'y'), this function returns a
    permutation of this list:

        [{'x': 0, 'y': 0},
        {'x': 0, 'y': 1},
        {'x': 1, 'y': 0},
        {'x': 1, 'y': 1}].

    Note that device IDs and device locations are equivalent. The former is a
    linearization of the latter along mesh dimensions.

Args:
        mesh: A TPU mesh.
        client_id: Optional; A DTensor client ID. If empty, query this client.


def events_from_logdir(logdir):
    assert tf.io.gfile.exists(logdir)
    files = tf.io.gfile.listdir(logdir)
    assert len(files) == 1, 'Found not exactly one file in logdir: %s' % files
    return _events_from_file(os.path.join(logdir, files[0]))

Returns all events in the single eventfile in logdir.

Args:
        logdir: The directory in which the single event file is sought.

Returns:
        A list of all tf.compat.v1.Event protos from the single event file.

Raises:
        AssertionError: If logdir does not contain exactly one file.


def __init__(self, config: AriaTextConfig, layer_idx: int):
    super().__init__(self)
    self.mlp = AriaTextMoELayer(config)

Aria Text Decoder Layer.

    This class defines a single decoder layer in the language model, incorporating self-attention and Mixture of Experts (MoE) feed-forward network.

Args:
        config (`AriaTextConfig`):
            Configuration object for the text component of the model.
        layer_idx (`int`):
            Index of the layer.


def load_stereo_chemical_props() -> Tuple[Mapping[str, List[Bond]], Mapping[str, List[Bond]], Mapping[str, List[BondAngle]]]:
    stereo_chemical_props = resources.read_text('openfold.resources', 'stereo_chemical_props.txt')
    lines_iter = iter(stereo_chemical_props.splitlines())
    residue_bonds: Dict[str, List[Bond]] = {}
    next(lines_iter)
    for line in lines_iter:
        if line.strip() == '-':
            break
        bond, resname, bond_length, stddev = line.split()
        atom1, atom2 = bond.split('-')
        if resname not in residue_bonds:
            residue_bonds[resname] = []
        residue_bonds[resname].append(Bond(atom1, atom2, float(bond_length), float(stddev)))
    residue_bonds['UNK'] = []
    residue_bond_angles: Dict[str, List[BondAngle]] = {}
    next(lines_iter)
    next(lines_iter)
    for line in lines_iter:
        if line.strip() == '-':
            break
        bond, resname, angle_degree, stddev_degree = line.split()
        atom1, atom2, atom3 = bond.split('-')
        if resname not in residue_bond_angles:
            residue_bond_angles[resname] = []
        residue_bond_angles[resname].append(BondAngle(atom1, atom2, atom3, float(angle_degree) / 180.0 * np.pi, float(stddev_degree) / 180.0 * np.pi))
    residue_bond_angles['UNK'] = []

    def make_bond_key(atom1_name: str, atom2_name: str) -> str:
        """Unique key to lookup bonds."""
        return '-'.join(sorted([atom1_name, atom2_name]))
    residue_virtual_bonds: Dict[str, List[Bond]] = {}
    for resname, bond_angles in residue_bond_angles.items():
        bond_cache: Dict[str, Bond] = {}
        for b in residue_bonds[resname]:
            bond_cache[make_bond_key(b.atom1_name, b.atom2_name)] = b
        residue_virtual_bonds[resname] = []
        for ba in bond_angles:
            bond1 = bond_cache[make_bond_key(ba.atom1_name, ba.atom2_name)]
            bond2 = bond_cache[make_bond_key(ba.atom2_name, ba.atom3name)]
            gamma = ba.angle_rad
            length = np.sqrt(bond1.length ** 2 + bond2.length ** 2 - 2 * bond1.length * bond2.length * np.cos(gamma))
            dl_outer = 0.5 / length
            dl_dgamma = 2 * bond1.length * bond2.length * np.sin(gamma) * dl_outer
            dl_db1 = (2 * bond1.length - 2 * bond2.length * np.cos(gamma)) * dl_outer
            dl_db2 = (2 * bond2.length - 2 * bond1.length * np.cos(gamma)) * dl_outer
            stddev = np.sqrt((dl_dgamma * ba.stddev) ** 2 + (dl_db1 * bond1.stddev) ** 2 + (dl_db2 * bond2.stddev) ** 2)
            residue_virtual_bonds[resname].append(Bond(ba.atom1_name, ba.atom3name, length, stddev))
    return (residue_bonds, residue_virtual_bonds, residue_bond_angles)

Load stereo_chemical_props.txt into a nice structure.

    Load literature values for bond lengths and bond angles and translate bond angles into the length of the opposite
    edge of the triangle ("residue_virtual_bonds").

Returns:
        residue_bonds: dict that maps resname --> list of Bond tuples residue_virtual_bonds: dict that maps resname -->
        list of Bond tuples residue_bond_angles: dict that maps resname --> list of BondAngle tuples


def diff_is_docstring_only(repo: Repo, branching_point: str, filename: str) -> bool:
    folder = Path(repo.working_dir)
    with checkout_commit(repo, branching_point):
        with open(folder / filename, 'r', encoding='utf-8') as f:
            old_content = f.read()
    with open(folder / filename, 'r', encoding='utf-8') as f:
        new_content = f.read()
    old_content_clean = clean_code(old_content)
    new_content_clean = clean_code(new_content)
    return old_content_clean == new_content_clean

Check if the diff is only in docstrings (or comments and whitespace) in a filename.

Args:
        repo (`git.Repo`): A git repository (for instance the Transformers repo).
        branching_point (`str`): The commit reference of where to compare for the diff.
        filename (`str`): The filename where we want to know if the diff isonly in docstrings/comments.

Returns:
        `bool`: Whether the diff is docstring/comments only or not.


def _unpack_sequence(self, state, n_before, n_after=-1):
    assert n_after >= -1
    state, seq = state.pop()
    options = []
    nontuple_seq = self.ctx.program.NewVariable()
    has_slurp = n_after > -1
    count = n_before + max(n_after, 0)
    nondeterministic_iterable = False
    for b in abstract_utils.expand_type_parameter_instances(seq.bindings):
        if b.data.full_name in ('builtins.set', 'builtins.frozenset'):
            nondeterministic_iterable = True
        tup = self._get_literal_sequence(b.data)
        if tup is not None:
            if has_slurp and len(tup) >= count:
                options.append(self._restructure_tuple(state, tup, n_before, n_after))
                continue
            elif len(tup) == count:
                options.append(tup)
                continue
            else:
                self.ctx.errorlog.bad_unpacking(self.frames, len(tup), count)
        if b.IsVisible(state.node):
            nontuple_seq.PasteBinding(b, state.node)
    if nontuple_seq.bindings:
        state, itr = self._get_iter(state, nontuple_seq)
        state, itr_result = self._call(state, itr, '__next__', ())
    elif not options:
        itr_result = self.ctx.new_unsolvable(state.node)
    else:
        itr_result = None
    if itr_result:
        option = [itr_result for _ in range(count)]
        if has_slurp:
            slurp = self.ctx.convert.build_list_of_type(state.node, itr_result)
            option = option[:n_before] + [slurp] + option[n_before:]
        options.append(option)
    values = tuple((self.ctx.convert.build_content(value, discard_concrete_values=False) for value in zip(*options)))
    if len(values) > 1 and nondeterministic_iterable:
        self.ctx.errorlog.nondeterministic_unpacking(self.frames)
    for value in reversed(values):
        if not value.bindings:
            value = self.ctx.convert.empty.to_variable(state.node)
        state = state.push(value)
    return state

Pops a tuple (or other iterable) and pushes it onto the VM's stack.

    Supports destructuring assignment with potentially a single list variable
    that slurps up the remaining elements:
    1. a, b, c = ...  # UNPACK_SEQUENCE
    2. a, *b, c = ... # UNPACK_EX

Args:
        state: The current VM state
        n_before: Number of elements before the list (n_elements for case 1)
        n_after: Number of elements after the list (-1 for case 1)

Returns:
        The new state.


def enable_save_as_bf16(variables: List[tf_variables.Variable]):
    for v in variables:
        if isinstance(v, d_variable.DVariable):
            v.save_as_bf16 = True

Allows float32 DVariables to be checkpointed and restored as bfloat16.

    The method only affects the DVariable part inside the model and leaves
    non-DTensor Variables/Tensors untouched.

Args:
        variables: A list of tf.Variable to be enabled with bfloat16 save/restore.
        Only has effect on DTensor Variables as they go through d_variables with
        DTensor Specific logis.


def split(state, num):
    state = tf_np.asarray(state, dtype=_RNG_KEY_DTYPE)
    state = _key2seed(state)
    try:
        states = stateless_random_ops.stateless_split(state, num)
    except AttributeError as e:
        states = stateless_split(state, num)
    states = array_ops_stack.unstack(states, num)
    states = nest.map_structure(_seed2key, states)
    return states

Creates new independent RNG states from an existing state.

Args:
        state: the existing state.
        num: the number of the new states.

Returns:
        A tuple of new states.


def is_initialized(self, name=None):
    if values_util.is_saving_non_distributed():
        return self._primary.is_initialized()
    if self._use_packed_variable():
        return self._packed_var.is_initialized()
    result = self._primary.is_initialized()
    for v in self._values[1:-1]:
        result = math_ops.logical_and(result, v.is_initialized())
    result = math_ops.logical_and(result, self._values[-1].is_initialized(), name=name)
    return result

Identifies if all the component variables are initialized.

Args:
        name: Name of the final `logical_and` op.

Returns:
        The op that evaluates to True or False depending on if all the
        component variables are initialized.


def watch_key(self):
    return _get_tensor_watch_key(self.node_name, self.output_slot, self.debug_op)

Watch key identities a debug watch on a tensor.

Returns:
        (`str`) A watch key, in the form of `tensor_name`:`debug_op`.


def forward(self, data: torch.Tensor, observed_indicator: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    ts_sum = (data * observed_indicator).abs().sum(self.dim, keepdim=True)
    num_observed = observed_indicator.sum(self.dim, keepdim=True)
    scale = ts_sum / torch.clamp(num_observed, min=1)
    if self.default_scale is None:
        batch_sum = ts_sum.sum(dim=0)
        batch_observations = torch.clamp(num_observed.sum(0), min=1)
        default_scale = torch.squeeze(batch_sum / batch_observations)
    else:
        default_scale = self.default_scale * torch.ones_like(scale)
    scale = torch.where(num_observed > 0, scale, default_scale)
    scale = torch.clamp(scale, min=self.minimum_scale)
    scaled_data = data / scale
    if not self.keepdim:
        scale = scale.squeeze(dim=self.dim)
    return (scaled_data, torch.zeros_like(scale), scale)

Parameters:
    data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`):
        input for Batch norm calculation
    observed_indicator (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`):
        Calculating the scale on the observed indicator.

Returns:
    tuple of `torch.Tensor` of shapes
        (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`,
        `(batch_size, 1, num_input_channels)`)


def _convert_from_saved_model(self, graph_def):
    self._save_conversion_params_metric(graph_def)
    quant_mode = QuantizationMode(self.optimizations, self.target_spec, self.representative_dataset, graph_def, self._experimental_disable_per_channel, self.experimental_new_dynamic_range_quantizer, self._experimental_low_bit_qat, self._experimental_full_integer_quantization_bias_type, self._experimental_variable_quantization, self._experimental_strict_qdq)
    self._validate_inference_input_output_types(quant_mode)
    converter_kwargs = {'enable_tflite_resource_variables': self.experimental_enable_resource_variables}
    converter_kwargs.update(self._get_base_converter_args())
    converter_kwargs.update(quant_mode.converter_flags())
    result = _convert_saved_model(**converter_kwargs)
    return self._optimize_tflite_model(result, quant_mode, _build_conversion_flags(**converter_kwargs).debug_options, quant_io=self.experimental_new_quantizer)

Helper method that converts saved model.

Args:
        graph_def: GraphDef object for the model, used only for stats.

Returns:
        The converted TFLite model.


def _list_valid_filenames_in_directory(directory, white_list_formats, split, class_indices, follow_links):
    dirname = os.path.basename(directory)
    if split:
        all_files = list(_iter_valid_files(directory, white_list_formats, follow_links))
        num_files = len(all_files)
        start, stop = (int(split[0] * num_files), int(split[1] * num_files))
        valid_files = all_files[start:stop]
    else:
        valid_files = _iter_valid_files(directory, white_list_formats, follow_links)
    classes = []
    filenames = []
    for root, fname in valid_files:
        classes.append(class_indices[dirname])
        absolute_path = os.path.join(root, fname)
        relative_path = os.path.join(dirname, os.path.relpath(absolute_path, directory))
        filenames.append(relative_path)
    return (classes, filenames)

Lists paths of files in `subdir` with extensions in `white_list_formats`.

Args:
        directory: absolute path to a directory containing the files to list.
            The directory name is used as class label
            and must be a key of `class_indices`.
        white_list_formats: set of strings containing allowed extensions for
            the files to be counted.
        split: tuple of floats (e.g. `(0.2, 0.6)`) to only take into
            account a certain fraction of files in each directory.
            E.g.: `segment=(0.6, 1.0)` would only account for last 40 percent
            of images in each directory.
        class_indices: dictionary mapping a class name to its index.
        follow_links: boolean, follow symbolic links to subdirectories.

Returns:
        classes: a list of class indices
        filenames: the path of valid files in `directory`, relative from
            `directory`'s parent (e.g., if `directory` is "dataset/class1",
            the filenames will be
            `["class1/file1.jpg", "class1/file2.jpg", ...]`).


def cast(cls, x, dtype):
    return x.astype(dtype)

Cast a tensor to a different dtype.

    Only called on a full array as provided by the user.

Args:
        x: the tensor to cast.
    Returns: the cast tensor.


def get_relevant_paths_and_versions(self, config: 'XLAConfigOptions'):
    if self.ld_library_path is None:
        self.ld_library_path = os.environ.get('LD_LIBRARY_PATH', None)
    if config.host_compiler == HostCompiler.CLANG:
        self.clang_path = _find_executable_or_die('clang', self.clang_path)
        self.clang_major_version = self.clang_major_version or _get_clang_major_version(self.clang_path)
        self.lld_path = self.lld_path or shutil.which('ld.lld')
    elif config.host_compiler == HostCompiler.GCC:
        self.gcc_path = _find_executable_or_die('gcc', self.gcc_path)
        self.gcc_major_version = self.gcc_major_version or _get_gcc_major_version(self.gcc_path)
    if config.backend == Backend.CUDA:
        if config.cuda_compiler == CudaCompiler.CLANG:
            self.clang_path = _find_executable_or_die('clang', self.clang_path)
        if not self.cuda_compute_capabilities:
            self.cuda_compute_capabilities = _get_cuda_compute_capabilities_or_die()

Gets paths and versions as needed by the config.

Args:
        config: XLAConfigOptions instance that determines what paths and versions
        to try to autoconfigure.


def visit_ImportFrom(self, node):
    if not node.module:
        self.generic_visit(node)
        return
    from_import = node.module
    from_import_first_component = from_import.split('.')[0]
    import_renames = getattr(self._api_change_spec, 'import_renames', {})
    import_rename_spec = import_renames.get(from_import_first_component, None)
    if not import_rename_spec:
        self.generic_visit(node)
        return
    updated_aliases = []
    same_aliases = []
    for import_alias in node.names:
        full_module_name = '%s.%s' % (from_import, import_alias.name)
        if excluded_from_module_rename(full_module_name, import_rename_spec):
            same_aliases.append(import_alias)
        else:
            updated_aliases.append(import_alias)
    if not updated_aliases:
        self.generic_visit(node)
        return
    assert self._stack[-1] is node
    parent = self._stack[-2]
    new_from_import = import_rename_spec.new_name + from_import[len(from_import_first_component):]
    updated_node = ast.ImportFrom(new_from_import, updated_aliases, node.level)
    ast.copy_location(updated_node, node)
    pasta.ast_utils.replace_child(parent, node, updated_node)
    additional_import_log = ''
    if same_aliases:
        same_node = ast.ImportFrom(from_import, same_aliases, node.level, col_offset=node.col_offset, lineno=node.lineno)
        ast.copy_location(same_node, node)
        parent.body.insert(parent.body.index(updated_node), same_node)
        pasta.base.formatting.set(same_node, 'prefix', pasta.base.formatting.get(updated_node, 'prefix'))
        additional_import_log = ' and %r' % pasta.dump(same_node)
    self.add_log(INFO, node.lineno, node.col_offset, 'Changed import from %r to %r%s.' % (pasta.dump(node), pasta.dump(updated_node), additional_import_log))
    self.generic_visit(node)

Handle visiting an import-from node in the AST.

Args:
        node: Current Node


def __init__(self, name='categorical_accuracy', dtype=None):
    super(CategoricalAccuracy, self).__init__(categorical_accuracy, name, dtype=dtype)

Calculates how often predictions match one-hot labels.

    You can provide logits of classes as `y_pred`, since argmax of
    logits and probabilities are same.

    This metric creates two local variables, `total` and `count` that are used to
    compute the frequency with which `y_pred` matches `y_true`. This frequency is
    ultimately returned as `categorical accuracy`: an idempotent operation that
    simply divides `total` by `count`.

    `y_pred` and `y_true` should be passed in as vectors of probabilities, rather
    than as labels. If necessary, use `tf.one_hot` to expand `y_true` as a vector.

    If `sample_weight` is `None`, weights default to 1.
    Use `sample_weight` of 0 to mask values.

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

    Standalone usage:

    >>> m = tf.keras.metrics.CategoricalAccuracy()
    >>> m.update_state([[0, 0, 1], [0, 1, 0]], [[0.1, 0.9, 0.8],
    ...                 [0.05, 0.95, 0]])
    >>> m.result().numpy()
    0.5

    >>> m.reset_state()
    >>> m.update_state([[0, 0, 1], [0, 1, 0]], [[0.1, 0.9, 0.8],
    ...                 [0.05, 0.95, 0]],
    ...                sample_weight=[0.7, 0.3])
    >>> m.result().numpy()
    0.3

    Usage with `compile()` API:

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


def _replace_tensors_by_numpy_ndarrays(repr_ds_map: rd.RepresentativeDatasetMapping) -> None:
    with session.Session() as sess:
        for signature_def_key in repr_ds_map:
            ds = repr_ds_map[signature_def_key]
            repr_ds_map[signature_def_key] = rd.replace_tensors_by_numpy_ndarrays(ds, sess)

Replaces tf.Tensors by their evaluated numpy arrays.

    This assumes that tf.Tensors in representative samples are created in the
    default Graph. It will raise an error if tensors are created in a different
    graph.

Args:
        repr_ds_map: SignatureDef key -> RepresentativeDataset mapping.


def _runOneBenchmark(self, default_device, num_iters=10, static_unroll=False, steps=10):

    def loop_body(i, x):
        with ops.device('/gpu:0'):
            nx = nn_ops.conv2d(input=x, filter=kernel, strides=[1, 1, 1, 1], padding='SAME', data_format='NHWC', name='conv2d')
            ni = math_ops.add(i, 1)
            return (ni, nx)
    ops.reset_default_graph()
    with session.Session() as sess, ops.device(default_device):
        i, x, kernel = self._getInitVariables()
        self.evaluate(variables.global_variables_initializer())
        if static_unroll:
            for _ in range(steps):
                i, x = loop_body(i, x)
        else:
            i, x = while_loop_tf.while_loop(lambda i, _: i < steps, loop_body, [i, x], parallel_iterations=steps, swap_memory=True)
        r = math_ops.reduce_sum(x)
        dx, dk = gradients_impl.gradients(r, [x, kernel])
        r = control_flow_ops.group(dx, dk)
        for _ in range(3):
            self.evaluate(r)
        start_time = time.time()
        for _ in range(num_iters):
            self.evaluate(r)
        return (time.time() - start_time) / num_iters

Evaluate the while loop performance.

Args:
        default_device: The default device to run all ops except the loop_body.
        loop_body is always run on GPU.
        num_iters: Number of iterations to run.
        static_unroll: If true, run unrolled version; otherwise, run while_loop.
        steps: Total number of repeated steps to run the loop.

Returns:
        The duration of the run in seconds.


def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
    mask = tf.cast(input_ids != padding_idx, tf.int64)
    incremental_indices = (tf.cumsum(mask, axis=1) + past_key_values_length) * mask
    return incremental_indices + padding_idx

Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
    are ignored. This is modified from fairseq's `utils.make_positions`.

Args:
        x: tf.Tensor x:

    Returns: tf.Tensor


def get_compiler_ir(self, device_name, platform_name, function_name, flat_args, captured_inputs, stage='hlo'):
    return pywrap_tfe.TF_GetCompilerIr(self._context_handle, function_name, stage, device_name, flat_args, captured_inputs, platform_name)

Get the compiler IR bytes.

Args:
        device_name: The name of the device with the form as
        "/job:localhost/replica:0/task:0/device:CPU:0", "/device:TPU:0" etc.
        When this is used, actual device is needed for getting the compiler IR.
        platform_name: The name of the platform, e.g. "TPU". When this is used,
        first we find a device whose name contains the platform, if it is found
        we get the compiler IR by device. Otherwise the compiler IR is obtained
        as if using that device. The former logic of falling back to device is
        necessary, as there are cases of TF variables that need to access
        devices, but the upper layer may generally choose platform for getting
        compiler IR in a device-agnostic way.
        function_name: The name of the function to get the compiler IR.
        flat_args: The flat argument inputs.
        captured_inputs: The inputs that are captured.
        stage: The exported stage for the given function.

Returns:
        The compiler IR bytes.


def _transpose_batch_time(x):
    x_static_shape = x.get_shape()
    if x_static_shape.rank is not None and x_static_shape.rank < 2:
        return x
    x_rank = array_ops.rank(x)
    x_t = array_ops.transpose(x, array_ops.concat(([1, 0], math_ops.range(2, x_rank)), axis=0))
    x_t.set_shape(tensor_shape.TensorShape([x_static_shape.dims[1].value, x_static_shape.dims[0].value]).concatenate(x_static_shape[2:]))
    return x_t

Transposes the batch and time dimensions of a Tensor.

    If the input tensor has rank < 2 it returns the original tensor. Retains as
    much of the static shape information as possible.

Args:
        x: A Tensor.

Returns:
        x transposed along the first two dimensions.


def _get(self, feed_item):
    result = store.get(self._entity, feed_item.get(FieldMap.CREATIVE_ASSET_ID, None))
    if not result:
        result = {'id': feed_item.get(FieldMap.CREATIVE_ASSET_ID, None), 'assetIdentifier': {'name': feed_item.get(FieldMap.CREATIVE_ASSET_NAME, None), 'type': feed_item.get(FieldMap.CREATIVE_TYPE, None)}}
        store.set(self._entity, [feed_item.get(FieldMap.CREATIVE_ASSET_ID, None)], result)
    return result

Retrieves an item from DCM or the local cache.

Args:
        feed_item: The feed item representing the creative asset from the
        Bulkdozer feed.

Returns:
        Instance of the DCM object either from the API or from the local cache.


def WriteUpdateTimestamp(self, update_timestamp=None):
    self.update_time = None
    if update_timestamp is None:
        update_timestamp = self._GetCurrentTime()
    return self._WriteTimestamp(update_timestamp, self.update_file)

Convenience method for writing the last update timestamp.

Args:
        update_timestamp: An int with the number of seconds since epoch,
        defaulting to the current time if None.

Returns:
        A boolean indicating success of the write.


def query_op_traceback(self, op_name):
    for op_log_proto in self._graph_tracebacks:
        for log_entry in op_log_proto.log_entries:
            if log_entry.name == op_name:
                return self._code_def_to_traceback(log_entry.code_def, op_log_proto.id_to_string)
    raise ValueError("Op '%s' does not exist in the tracebacks received by the debug server." % op_name)

Query the traceback of an op.

Args:
        op_name: Name of the op to query.

Returns:
        The traceback of the op, as a list of 3-tuples:
        (filename, lineno, function_name)

Raises:
        ValueError: If the op cannot be found in the tracebacks received by the
        server so far.


def std(x, axis=None, keepdims=False):
    if any_symbolic_tensors((x,)):
        return Std(axis=axis, keepdims=keepdims).symbolic_call(x)
    return backend.numpy.std(x, axis=axis, keepdims=keepdims)

Compute the standard deviation along the specified axis.

Args:
        x: Input tensor.
        axis: Axis along which to compute standard deviation.
            Default is to compute the standard deviation of the
            flattened tensor.
        keepdims: If this is set to `True`, the axes which are reduced are left
            in the result as dimensions with size one.

Returns:
        Output tensor containing the standard deviation values.


def update_with_token(self, token_id: int) -> bool:
    if self.status != RequestStatus.DECODING:
        return False
    is_eos = token_id == self.eos_token_id and self.eos_token_id != -1
    is_max_len = self.generated_len() >= self.max_new_tokens
    if is_eos or is_max_len:
        self.status = RequestStatus.FINISHED
        return True
    return False

Update the request with a newly generated token and check for completion.

Args:
        token_id: The token ID to add to the output sequence

Returns:
        bool: True if the request is now complete, False otherwise


def _unsorted_segment_N(data, segment_ids, num_segments):
    num_segments = ops.convert_to_tensor(num_segments)
    segment_ids_shape = array_ops.shape_internal(segment_ids)
    ones_tensor = array_ops.ones(segment_ids_shape, dtype=data.dtype)
    n = gen_math_ops.unsorted_segment_sum(ones_tensor, segment_ids, num_segments)
    broadcastable_shape = array_ops.concat([num_segments[array_ops.newaxis], array_ops.ones([array_ops.rank(data) - array_ops.rank(segment_ids)], dtype=num_segments.dtype)], axis=0)
    n = array_ops.reshape(n, broadcastable_shape)
    return gen_math_ops.maximum(n, 1)

Helper function for unsorted_segment_mean/_sqrtN.

    Computes the number of segment entries with 0-entries set to 1 to allow
    division by N.

Args:
        data: A `Tensor` with data that will be assembled in the output.
        segment_ids: An integer tensor whose shape is a prefix of `data.shape`. The
        values must be in the range `[0, num_segments)`. The values are always
        validated to be in range on CPU, never validated on TPU/GPU.
        num_segments: An integer scalar `Tensor`. The number of distinct segment
        IDs.

Returns:
        A `Tensor` with the number of segment entries with 0-entries set to 1.


Callback invoked during construction of the debug-wrapper session.

    This is a blocking callback.
    The invocation happens right before the constructor ends.

Args:
        request: (`OnSessionInitRequest`) callback request carrying information
        such as the session being wrapped.

Returns:
        An instance of `OnSessionInitResponse`.


def make_pixel_mask(image: np.ndarray, output_size: Tuple[int, int], input_data_format: Optional[Union[str, ChannelDimension]]=None) -> np.ndarray:
    input_height, input_width = get_image_size(image, channel_dim=input_data_format)
    mask = np.zeros(output_size, dtype=np.int64)
    mask[:input_height, :input_width] = 1
    return mask

Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding.

Args:
        image (`np.ndarray`):
            Image to make the pixel mask for.
        output_size (`Tuple[int, int]`):
            Output size of the mask.


def adapt(self, data, batch_size=None, steps=None):
    self.reset_state()
    if isinstance(data, tf.data.Dataset):
        if steps is not None:
            data = data.take(steps)
        for batch in data:
            self.update_state(batch)
    else:
        data = tf_utils.ensure_tensor(data, dtype='string')
        if data.shape.rank == 1:
            data = tf.expand_dims(data, -1)
        self.update_state(data)
    self.finalize_state()

Computes a vocabulary of string terms from tokens in a dataset.

    Calling `adapt()` on a `TextVectorization` layer is an alternative to
    passing in a precomputed vocabulary on construction via the `vocabulary`
    argument. A `TextVectorization` layer should always be either adapted
    over a dataset or supplied with a vocabulary.

    During `adapt()`, the layer will build a vocabulary of all string tokens
    seen in the dataset, sorted by occurrence count, with ties broken by
    sort order of the tokens (high to low). At the end of `adapt()`, if
    `max_tokens` is set, the vocabulary will be truncated to `max_tokens`
    size. For example, adapting a layer with `max_tokens=1000` will compute
    the 1000 most frequent tokens occurring in the input dataset. If
    `output_mode='tf-idf'`, `adapt()` will also learn the document
    frequencies of each token in the input dataset.

Args:
        data: The data to train on. It can be passed either as a
            batched `tf.data.Dataset`, as a list of strings,
            or as a NumPy array.
        steps: Integer or `None`.
            Total number of steps (batches of samples) to process.
            If `data` is a `tf.data.Dataset`, and `steps` is `None`,
            `adapt()` will run until the input dataset is exhausted.
            When passing an infinitely
            repeating dataset, you must specify the `steps` argument. This
            argument is not supported with array inputs or list inputs.


def is_ipython_subprocess() -> bool:
    return False

Check if we are in a sub-process launched from within a `ipython` terminal.

Returns:
        `True` only if we are in ipython terminal (e.g. `ml_python`) and inside
        a sub-process.


def required_space_to_batch_paddings(input_shape, block_shape, base_paddings=None, name=None):
    with ops.name_scope(name, 'required_space_to_batch_paddings', [input_shape, block_shape]):
        input_shape = ops.convert_to_tensor(input_shape, dtype=dtypes.int32, name='input_shape')
        block_shape = ops.convert_to_tensor(block_shape, dtype=dtypes.int32, name='block_shape')
        block_shape.get_shape().assert_is_fully_defined()
        block_shape.get_shape().assert_has_rank(1)
        num_block_dims = block_shape.get_shape().dims[0].value
        if num_block_dims == 0:
            return (zeros([0, 2], dtypes.int32), zeros([0, 2], dtypes.int32))
        input_shape.get_shape().assert_is_compatible_with([num_block_dims])
        if base_paddings is not None:
            base_paddings = ops.convert_to_tensor(base_paddings, dtype=dtypes.int32, name='base_paddings')
            base_paddings.get_shape().assert_is_compatible_with([num_block_dims, 2])
        else:
            base_paddings = zeros([num_block_dims, 2], dtypes.int32)
        const_block_shape = tensor_util.constant_value(block_shape)
        const_input_shape = tensor_util.constant_value(input_shape)
        const_base_paddings = tensor_util.constant_value(base_paddings)
        if const_block_shape is not None and const_input_shape is not None and (const_base_paddings is not None):
            block_shape = const_block_shape
            input_shape = const_input_shape
            base_paddings = const_base_paddings
        pad_start = base_paddings[:, 0]
        orig_pad_end = base_paddings[:, 1]
        full_input_shape = input_shape + pad_start + orig_pad_end
        pad_end_extra = (block_shape - full_input_shape % block_shape) % block_shape
        pad_end = orig_pad_end + pad_end_extra
        result_paddings = array_ops_stack.stack([[pad_start[i], pad_end[i]] for i in range(num_block_dims)], name='paddings')
        result_crops = array_ops_stack.stack([[0, pad_end_extra[i]] for i in range(num_block_dims)], name='crops')
        return (result_paddings, result_crops)

Calculate padding required to make block_shape divide input_shape.

    This function can be used to calculate a suitable paddings argument for use
    with space_to_batch_nd and batch_to_space_nd.

Args:
        input_shape: int32 Tensor of shape [N].
        block_shape: int32 Tensor of shape [N].
        base_paddings: Optional int32 Tensor of shape [N, 2].  Specifies the minimum
        amount of padding to use.  All elements must be >= 0.  If not specified,
        defaults to 0.
        name: string.  Optional name prefix.

Returns:
        (paddings, crops), where:

        `paddings` and `crops` are int32 Tensors of rank 2 and shape [N, 2]
        satisfying:

        paddings[i, 0] = base_paddings[i, 0].
        0 <= paddings[i, 1] - base_paddings[i, 1] < block_shape[i]
        (input_shape[i] + paddings[i, 0] + paddings[i, 1]) % block_shape[i] == 0

        crops[i, 0] = 0
        crops[i, 1] = paddings[i, 1] - base_paddings[i, 1]

    Raises: ValueError if called with incompatible shapes.


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

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

Tests for errors in the structured signature.

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


def get_dofn_specs(dofn: 'DoFn') -> tuple[set[StateSpec], set[TimerSpec]]:
    from apache_beam.runners.common import MethodWrapper
    from apache_beam.transforms.core import _DoFnParam
    from apache_beam.transforms.core import _StateDoFnParam
    from apache_beam.transforms.core import _TimerDoFnParam
    all_state_specs = set()
    all_timer_specs = set()
    for method_name in dir(dofn):
        if not isinstance(getattr(dofn, method_name, None), types.MethodType):
            continue
        method = MethodWrapper(dofn, method_name)
        param_ids = [d.param_id for d in method.defaults if isinstance(d, _DoFnParam)]
        if len(param_ids) != len(set(param_ids)):
            raise ValueError('DoFn %r has duplicate %s method parameters: %s.' % (dofn, method_name, param_ids))
        for d in method.defaults:
            if isinstance(d, _StateDoFnParam):
                all_state_specs.add(d.state_spec)
            elif isinstance(d, _TimerDoFnParam):
                all_timer_specs.add(d.timer_spec)
    return (all_state_specs, all_timer_specs)

Gets the state and timer specs for a DoFn, if any.

Args:
        dofn (apache_beam.transforms.core.DoFn): The DoFn instance to introspect for
        timer and state specs.


def ndim(x):
    if any_symbolic_tensors((x,)):
        return Ndim().symbolic_call(x)
    return backend.numpy.ndim(x)

Return the number of dimensions of a tensor.

Args:
        x: Input tensor.

Returns:
        The number of dimensions in `x`.


def f(options, expected_tf_failures=0):
    test_parameters = [{'ksize': [[2, 1, 1, 2], [1, 1, 1, 1], [1, 1, 2, 1], [1, 10, 11, 1]], 'strides': [[2, 1, 1, 2], [1, 1, 1, 1], [1, 1, 2, 1], [1, 10, 11, 1]], 'input_shape': [[], [1, 1, 1, 1], [1, 15, 14, 1], [3, 15, 14, 3]], 'padding': ['SAME', 'VALID'], 'data_format': ['NHWC'], 'fully_quantize': [False], 'quant_16x8': [False]}, {'ksize': [[2, 1, 1, 2], [1, 1, 1, 1], [1, 1, 2, 1], [1, 10, 11, 1]], 'strides': [[2, 1, 1, 2], [1, 1, 1, 1], [1, 1, 2, 1], [1, 10, 11, 1]], 'input_shape': [[], [1, 1, 1, 1], [1, 15, 14, 1], [3, 15, 14, 3]], 'padding': ['SAME', 'VALID'], 'data_format': ['NHWC'], 'fully_quantize': [True], 'quant_16x8': [False]}, {'ksize': [[1, 1, 1, 1]], 'strides': [[1, 1, 1, 1]], 'input_shape': [[1, 1, 1, 1]], 'padding': ['SAME', 'VALID'], 'data_format': ['NHWC'], 'fully_quantize': [True], 'quant_16x8': [True]}]
    if not allow_fully_quantize:
        test_parameters = [test_parameter for test_parameter in test_parameters if True not in test_parameter['fully_quantize']]

    def build_graph(parameters):
        input_tensor = tf.compat.v1.placeholder(dtype=tf.float32, name='input', shape=parameters['input_shape'])
        out = pool_op(input_tensor, ksize=parameters['ksize'], strides=parameters['strides'], data_format=parameters['data_format'], padding=parameters['padding'])
        return ([input_tensor], [out])

    def build_inputs(parameters, sess, inputs, outputs):
        if allow_fully_quantize:
            input_values = create_tensor_data(tf.float32, parameters['input_shape'], min_value=-1, max_value=1)
        else:
            input_values = create_tensor_data(tf.float32, parameters['input_shape'])
        return ([input_values], sess.run(outputs, feed_dict=dict(zip(inputs, [input_values]))))
    make_zip_of_tests(options, test_parameters, build_graph, build_inputs, expected_tf_failures=expected_tf_failures)

Actual function that generates examples.

Args:
        options: An Options instance.
        expected_tf_failures: number of expected tensorflow failures.


def download_artifacts_from_gcs(bucket_name, prefix, local_path):
    client = Client()
    bucket = client.get_bucket(bucket_name)
    blobs = [blob.name for blob in bucket.list_blobs(prefix=prefix)]
    _ = transfer_manager.download_many_to_path(bucket, blobs, destination_directory=local_path)

Downloads artifacts from GCS to the local file system.

Args:
        bucket_name: The name of the GCS bucket to download from.
        prefix: Prefix of GCS objects to download.
        local_path: The local path to download the folder to.


def fetch_layout(tensor: tensor_lib.Tensor) -> layout_lib.Layout:
    return _dtensor_device().fetch_layout(tensor)

Fetches the layout of a DTensor.

Args:
        tensor: The DTensor whose layout is to be fetched.

Returns:
        The `Layout` of this DTensor.

Raises:
        RuntimeError: When not called eagerly.


def GetUpdates(self, gcs_client, bucket_name, obj, since):
    bucket = gcs_client.bucket(bucket_name)
    blob = bucket.get_blob(obj)
    if blob is None:
        self.log.error('GCS object gs://%s/%s not found', bucket_name, obj)
        raise error.SourceUnavailable('unable to download object from GCS.')
    if since and timestamps.FromDateTimeToTimestamp(blob.updated) < since:
        return []
    data_map = self.GetMap(cache_info=blob.open())
    data_map.SetModifyTimestamp(timestamps.FromDateTimeToTimestamp(blob.updated))
    return data_map

Gets updates from a source.

Args:
        gcs_client: initialized gcs client
        bucket_name: gcs bucket name
        obj: object with the data
        since: a timestamp representing the last change (None to force-get)

Returns:
        A tuple containing the map of updates and a maximum timestamp