def run_inference(self, batch: Sequence[datatable.Frame], model: Union[xgboost.Booster, xgboost.XGBModel], inference_args: Optional[dict[str, Any]]=None) -> Iterable[PredictionResult]:
    return self._inference_fn(batch, model, inference_args)

Runs inferences on a batch of datatable dataframe.

Args:
        batch: A sequence of examples as datatable dataframes. Each
        row in a dataframe is a single example. The dimensions
        must match the dimensions of the data used to train
        the model.
        model: XGBoost booster or XBGModel (sklearn interface). Must implement
        predict(X). Where the parameter X is a datatable dataframe.
        inference_args: Any additional arguments for an inference.

Returns:
        An Iterable of type PredictionResult.


def assert_eventual(self, func, required, allowed, timeout_secs=300.0):
    required = set(required)
    assert required
    seen = set()
    start_time = time.time()
    while timeout_secs is None or time.time() - start_time < timeout_secs:
        if seen == required:
            return
        value = func()
        if value not in allowed:
            self.fail(msg=f'Disallowed value: {value}.')
        if value in required:
            seen.add(value)
    missing = [v for v in required if v not in seen]
    self.fail(msg=f'Timed out. Missing values: {str([str(v) for v in missing])}.')

Tests that calls to the given function meet required and allowed values.

Args:
        func: function to test.
        required: iterable of required values. Must be hashable and non-empty.
        allowed: iterable of allowed values. Must be hashable and non-empty.
        timeout_secs: fails if more than this time is required.


def _build_all_reduce_ring(core_locations: List[_CoreLocation], rotate: bool=False) -> List[int]:
    permutation = list(range(len(core_locations)))
    if not permutation:
        return permutation
    logging.vlog(2, 'Core locations in: %s', core_locations)
    first_column = min([l.x for l in core_locations])
    first_row = min([l.y for l in core_locations])
    same_z = len(set([l.z for l in core_locations])) == 1
    logging.vlog(2, 'first_column: %d', first_column)
    logging.vlog(2, 'first_row: %d', first_row)
    logging.vlog(2, 'same_z: %s', same_z)

    def _cmp_2d(ia: int, ib: int) -> int:
        if not rotate:
            a = core_locations[ia]
            b = core_locations[ib]
            a_first = a.x == first_column and a.y != first_row
            b_first = b.x == first_column and b.y != first_row
            if a_first != b_first:
                return -1 if b_first else 1
            if a.y != b.y:
                return b.y - a.y if a_first else a.y - b.y
            if a.x != b.x:
                return a.x - b.x if a.y % 2 == 0 else b.x - a.x
            return a.core - b.core
        else:
            a = core_locations[ia]
            b = core_locations[ib]
            a_first = a.y == first_row and a.x != first_column
            b_first = b.y == first_row and b.x != first_column
            if a_first != b_first:
                return -1 if b_first else 1
            if a.x != b.x:
                return b.x - a.x if a_first else a.x - b.x
            if a.y != b.y:
                return a.y - b.y if a.x % 2 == 0 else b.y - a.y
            return a.core - b.core

    def _cmp_3d(ia: int, ib: int) -> int:
        a = core_locations[ia]
        b = core_locations[ib]
        a_corner = a.x == first_column and a.y == first_row
        b_corner = b.x == first_column and b.y == first_row
        if a_corner and b_corner:
            return b.z - a.z if a.z != b.z else a.core - b.core
        if a_corner != b_corner:
            return -1 if b_corner else 1
        if a.z == b.z:
            return _cmp_2d(ia, ib) if a.z % 2 == 0 else -_cmp_2d(ia, ib)
        return a.z - b.z
    if same_z:
        permutation.sort(key=functools.cmp_to_key(_cmp_2d))
    else:
        permutation.sort(key=functools.cmp_to_key(_cmp_3d))
    logging.vlog(2, 'Permutation out: %s', permutation)
    return permutation

Reorders a list of TPU cores to optimize for AllReduce performance.

    This is ported from the C++ tensorflow::BuildAllReduceRing function,
    mixed with some logic from TF TPU's device_assignment._ring_3d.

Args:
        core_locations: A list of core locations expressed as [x, y, z, core].
        rotate: If true, scan the cores in a column-major order. False by default.

Returns:
        A permutation of the input list such that neighbors in the sequence are
        nearby in the TPU topology.


def update_state(self, y_true, y_pred, sample_weight=None):
    return metrics_utils.update_confusion_matrix_variables({self._confusion_matrix_cond: self.accumulator}, y_true, y_pred, thresholds=self.thresholds, thresholds_distributed_evenly=self._thresholds_distributed_evenly, sample_weight=sample_weight)

Accumulates the metric statistics.

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


def get_extended_attention_mask(self, attention_mask: torch.Tensor, input_shape: Tuple[int], device: torch.device, has_query: bool=False) -> torch.Tensor:
    if attention_mask.dim() == 3:
        extended_attention_mask = attention_mask[:, None, :, :]
    elif attention_mask.dim() == 2:
        extended_attention_mask = attention_mask[:, None, None, :]
    else:
        raise ValueError('Wrong shape for input_ids (shape {}) or attention_mask (shape {})'.format(input_shape, attention_mask.shape))
    extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)
    extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
    return extended_attention_mask

Makes broadcastable attention and causal masks so that future and masked tokens are ignored.

Args:
        attention_mask (`torch.Tensor`):
            Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
        input_shape (`Tuple[int]`):
            The shape of the input to the model.
        device (`torch.device`):
            The device of the input to the model.

Returns:
        `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`.


def decode(self, token_ids, **kwargs):
    return super().decode(token_ids, **kwargs)

Converts a sequence of ids in a string, using the tokenizer and vocabulary with options to remove special
    tokens and clean up tokenization spaces.

    Similar to doing `self.convert_tokens_to_string(self.convert_ids_to_tokens(token_ids))`.

Args:
        token_ids (`Union[int, List[int], np.ndarray, torch.Tensor, tf.Tensor]`):
            List of tokenized input ids. Can be obtained using the `__call__` method.
        skip_special_tokens (`bool`, *optional*, defaults to `False`):
            Whether or not to remove special tokens in the decoding.
        clean_up_tokenization_spaces (`bool`, *optional*):
            Whether or not to clean up the tokenization spaces. If `None`, will default to
            `self.clean_up_tokenization_spaces` (available in the `tokenizer_config`).
        use_source_tokenizer (`bool`, *optional*, defaults to `False`):
            Whether or not to use the source tokenizer to decode sequences (only applicable in sequence-to-sequence
            problems).
        kwargs (additional keyword arguments, *optional*):
            Will be passed to the underlying model specific decode method.

Returns:
        `str`: The decoded sentence.


def _normalize(self, text: str) -> str:
    accepted = [chr(i) for i in range(ord('a'), ord('z') + 1)] + [chr(i) for i in range(ord('A'), ord('Z') + 1)] + [chr(i) for i in range(ord('0'), ord('9') + 1)] + ['.']
    accepted = frozenset(accepted)
    pattern = re.compile('_+')
    text = ''.join([c if c in accepted else '_' for c in text.lower()])
    text = pattern.sub('_', text).strip('_')
    return text

Normalizes the input text. This process is for the genres and the artist

Args:
        text (`str`):
            Artist or Genre string to normalize


def forward(self, input_ids: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple, BaseModelOutputWithPooling]:
    return self.text_model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)

Examples:

    ```python
    >>> from transformers import AutoTokenizer, XCLIPTextModel

    >>> model = XCLIPTextModel.from_pretrained("microsoft/xclip-base-patch32")
    >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/xclip-base-patch32")

    >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")

    >>> outputs = model(**inputs)
    >>> last_hidden_state = outputs.last_hidden_state
    >>> pooled_output = outputs.pooler_output  # pooled (EOS token) states
    ```


def __init__(self, bits: int=4, p: int=2, modules_to_not_convert: Optional[List[str]]=None, hadamard_size: int=512, group_size: int=256, tune_metadata: Optional[Dict[str, Any]]=None, **kwargs):
    if tune_metadata is None:
        tune_metadata = {}
    self.quant_method = QuantizationMethod.HIGGS
    self.bits = bits
    self.p = p
    self.modules_to_not_convert = modules_to_not_convert
    self.hadamard_size = hadamard_size
    self.group_size = group_size
    self.tune_metadata = tune_metadata
    self.post_init()

HiggsConfig is a configuration class for quantization using the HIGGS method.

Args:
        bits (int, *optional*, defaults to 4):
            Number of bits to use for quantization. Can be 2, 3 or 4. Default is 4.
        p (int, *optional*, defaults to 2):
            Quantization grid dimension. 1 and 2 are supported. 2 is always better in practice. Default is 2.
        modules_to_not_convert (`list`, *optional*, default to ["lm_head"]):
            List of linear layers that should not be quantized.
        hadamard_size (int, *optional*, defaults to 512):
            Hadamard size for the HIGGS method. Default is 512. Input dimension of matrices is padded to this value. Decreasing this below 512 will reduce the quality of the quantization.
        group_size (int, *optional*, defaults to 256):
            Group size for the HIGGS method. Can be 64, 128 or 256. Decreasing it barely affects the performance. Default is 256. Must be a divisor of hadamard_size.
        tune_metadata ('dict', *optional*, defaults to {}):
            Module-wise metadata (gemm block shapes, GPU metadata, etc.) for saving the kernel tuning results. Default is an empty dictionary. Is set automatically during tuning.


def basic_train_loop(supervisor, train_step_fn, args=None, kwargs=None, master=''):
    if args is None:
        args = []
    if kwargs is None:
        kwargs = {}
    should_retry = True
    while should_retry:
        try:
            should_retry = False
            with supervisor.managed_session(master) as sess:
                while not supervisor.should_stop():
                    train_step_fn(sess, *args, **kwargs)
        except errors.AbortedError:
            should_retry = True

Basic loop to train a model.

    Calls `train_step_fn` in a loop to train a model.  The function is called as:

    ```python
    train_step_fn(session, *args, **kwargs)
    ```

    It is passed a `tf.compat.v1.Session` in addition to `args` and `kwargs`.  The
    function
    typically runs one training step in the session.

Args:
        supervisor: `tf.compat.v1.train.Supervisor` to run the training services.
        train_step_fn: Callable to execute one training step.  Called repeatedly as
        `train_step_fn(session, *args **kwargs)`.
        args: Optional positional arguments passed to `train_step_fn`.
        kwargs: Optional keyword arguments passed to `train_step_fn`.
        master: Master to use to create the training session.  Defaults to `""`
        which causes the session to be created in the local process.


def get_factors(dividend: int) -> Set[int]:
    factors_set = set()
    for i in range(1, int(dividend ** 0.5) + 1):
        if dividend % i == 0:
            factors_set.add(i)
            factors_set.add(dividend // i)
    return factors_set

Calculate all factors of a given number, i.e. a divisor that leaves
    no remainder. For example, if dividend=12, it will return {1, 2, 3, 4, 6, 12}.

Args:
        dividend (int): The number to find factors for.

Returns:
        set: A set containing all factors of the number.


def translate_distribution(distribution_update, metric_update_proto):
    dist_update_proto = dataflow.DistributionUpdate()
    dist_update_proto.min = to_split_int(distribution_update.min)
    dist_update_proto.max = to_split_int(distribution_update.max)
    dist_update_proto.count = to_split_int(distribution_update.count)
    dist_update_proto.sum = to_split_int(distribution_update.sum)
    if isinstance(distribution_update, DataflowDistributionCounter):
        dist_update_proto.histogram = dataflow.Histogram()
        distribution_update.translate_to_histogram(dist_update_proto.histogram)
    metric_update_proto.distribution = dist_update_proto

Translate metrics DistributionUpdate to dataflow distribution update.

Args:
        distribution_update: Instance of DistributionData,
        DistributionInt64Accumulator or DataflowDistributionCounter.
        metric_update_proto: Used for report metrics.


def __init__(self, layer, **kwargs):
    if not isinstance(layer, Layer):
        raise ValueError(f'Please initialize `TimeDistributed` layer with a `keras.layers.Layer` instance. Received: {layer}')
    super().__init__(layer, **kwargs)
    self.supports_masking = True

This wrapper allows to apply a layer to every temporal slice of an input.

    Every input should be at least 3D, and the dimension of index one of the
    first input will be considered to be the temporal dimension.

    Consider a batch of 32 video samples, where each sample is a 128x128 RGB
    image with `channels_last` data format, across 10 timesteps.
    The batch input shape is `(32, 10, 128, 128, 3)`.

    You can then use `TimeDistributed` to apply the same `Conv2D` layer to each
    of the 10 timesteps, independently:

    >>> inputs = layers.Input(shape=(10, 128, 128, 3), batch_size=32)
    >>> conv_2d_layer = layers.Conv2D(64, (3, 3))
    >>> outputs = layers.TimeDistributed(conv_2d_layer)(inputs)
    >>> outputs.shape
    (32, 10, 126, 126, 64)

    Because `TimeDistributed` applies the same instance of `Conv2D` to each of
    the timestamps, the same set of weights are used at each timestamp.

Args:
        layer: a `keras.layers.Layer` instance.

    Call arguments:
        inputs: Input tensor of shape (batch, time, ...) or nested tensors,
            and each of which has shape (batch, time, ...).
        training: Python boolean indicating whether the layer should behave in
            training mode or in inference mode. This argument is passed to the
            wrapped layer (only if the layer supports this argument).
        mask: Binary tensor of shape `(samples, timesteps)` indicating whether
            a given timestep should be masked. This argument is passed to the
            wrapped layer (only if the layer supports this argument).


def __init__(self, data_format=None, **kwargs):
    super().__init__(**kwargs)
    self.data_format = backend.standardize_data_format(data_format)
    self.input_spec = InputSpec(min_ndim=1)
    self._channels_first = self.data_format == 'channels_first'

Flattens the input. Does not affect the batch size.

    Note: If inputs are shaped `(batch,)` without a feature axis, then
    flattening adds an extra channel dimension and output shape is `(batch, 1)`.

Args:
        data_format: A string, one of `"channels_last"` (default) or
            `"channels_first"`. The ordering of the dimensions in the inputs.
            `"channels_last"` corresponds to inputs with shape
            `(batch, ..., channels)` while `"channels_first"` corresponds to
            inputs with shape `(batch, channels, ...)`.
            When unspecified, uses `image_data_format` value found in your Keras
            config file at `~/.keras/keras.json` (if exists). Defaults to
            `"channels_last"`.

Example:
    >>> x = keras.Input(shape=(10, 64))
    >>> y = keras.layers.Flatten()(x)
    >>> y.shape
    (None, 640)


def _load_method(self, state, self_obj, name):
    state, result = self.load_attr(state, self_obj, name)
    state = self._push_null(state)
    return (state.push(result), result)

Loads and pushes a method on the stack.

Args:
        state: the current VM state.
        self_obj: the `self` object of the method.
        name: the name of the method.

Returns:
        (state, method) where `state` is the updated VM state and `method` is the
        method that was loaded. The method is already pushed onto the stack,
        either at the top or below the `self` object.


def test_pass(self, e=None):
    self._test_end(TestResultEnums.TEST_RESULT_PASS, e)

To mark the test as passed in this record.

Args:
        e: An instance of mobly.signals.TestPass.


def split_if_relative_reference(reference: message.Message) -> None:
    _validate_reference(reference)
    uri_field = reference.DESCRIPTOR.fields_by_name.get('uri')
    if not proto_utils.field_is_set(reference, uri_field):
        return
    uri = proto_utils.get_value_at_field(reference, uri_field)
    internal_match = re.fullmatch(_INTERNAL_REFERENCE_PATTERN, uri.value)
    if internal_match is not None:
        reference_id_field = get_reference_id_field_for_resource(reference, internal_match.group('resource_type'))
        reference_id = proto_utils.create_message_from_descriptor(reference_id_field.message_type)
        populate_typed_reference_id(reference_id, internal_match.group('resource_id'), internal_match.group('version'))
        proto_utils.copy_common_field(uri, reference_id, 'id')
        proto_utils.copy_common_field(uri, reference_id, 'extension')
        proto_utils.set_value_at_field(reference, reference_id_field, reference_id)
        return
    fragment_match = re.fullmatch(_FRAGMENT_REFERENCE_PATTERN, uri.value)
    if fragment_match is not None:
        fragment_field = reference.DESCRIPTOR.fields_by_name['fragment']
        fragment = proto_utils.create_message_from_descriptor(fragment_field.message_type)
        value_field = fragment.DESCRIPTOR.fields_by_name['value']
        proto_utils.set_value_at_field(fragment, value_field, uri.value[1:])
        proto_utils.copy_common_field(uri, fragment, 'id')
        proto_utils.copy_common_field(uri, fragment, 'extension')
        proto_utils.set_value_at_field(reference, fragment_field, fragment)
        return

If possible, parses a `Reference` `uri` into more structured fields.

    This is only possible for two forms of reference uris:
    * Relative references of the form $TYPE/$ID, e.g., "Patient/1234"
        In this case, this will be parsed to a proto of the form:
        {patient_id: {value: "1234"}}
    * Fragments of the form "#$FRAGMENT", e.g., "#vs1".  In this case, this would
        be parsed into a proto of the form:
        {fragment: {value: "vs1"} }

    If the reference URI matches one of these schemas, the `uri` field will be
    cleared, and the appropriate structured fields set. Otherwise, the reference
    will be unchanged.

Args:
        reference: The FHIR reference to potentially split.

Raises:
        ValueError: If the message is not a valid FHIR Reference proto.


def get_container_image_from_options(pipeline_options):
    worker_options = pipeline_options.view_as(WorkerOptions)
    if worker_options.sdk_container_image:
        return worker_options.sdk_container_image
    container_repo = names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY
    image_name = '{repository}/beam_python{major}.{minor}_sdk'.format(repository=container_repo, major=sys.version_info[0], minor=sys.version_info[1])
    image_tag = _get_required_container_version()
    return image_name + ':' + image_tag

For internal use only; no backwards-compatibility guarantees.

Args:
        pipeline_options (PipelineOptions): A container for pipeline options.

Returns:
        str: Container image for remote execution.


def symmetric_kl_divergence(predicted, actual):
    epsilon = tf.constant(1e-07, dtype=tf.float32, name='epsilon')
    p = tf.math.maximum(predicted, epsilon)
    q = tf.math.maximum(actual, epsilon)
    kld_1 = tf.math.reduce_sum(tf.math.multiply(p, tf.math.log(tf.math.divide(p, q))))
    kld_2 = tf.math.reduce_sum(tf.math.multiply(q, tf.math.log(tf.math.divide(q, p))))
    return tf.add(kld_1, kld_2)

Calculate symmetric KL-divergence over two classification tensors.

    Note that here the classifications do not form a probability distribution.
    They are, however normalized to 0..1 and calculating a KL-divergence over them
    gives reasonable numerical results.

    Shape of the two inputs must be the same at inference time but is otherwise
    unconstrained.

Args:
        predicted: classification outputs from model
        actual: golden classification outputs

Returns:
        Single scalar tensor with symmetric KL-divergence between predicted and
        actual.


def __init__(self, value_type, default: typing.Optional[numbers.Number]=MISSING_VALUE, min_value: typing.Optional[numbers.Number]=None, max_value: typing.Optional[numbers.Number]=None, is_noneable: bool=False, frozen: bool=False):
    if min_value is not None and max_value is not None and (min_value > max_value):
        raise ValueError(f'"max_value" must be equal or greater than "min_value". Encountered: min_value={min_value}, max_value={max_value}.')
    self._min_value = min_value
    self._max_value = max_value
    super().__init__(value_type, default, is_noneable=is_noneable, frozen=frozen)

Constructor.

Args:
        value_type: Type of number.
        default: Default value for this spec.
        min_value: (Optional) minimum value of acceptable values.
        max_value: (Optional) maximum value of acceptable values.
        is_noneable: If True, None is acceptable.
        frozen: If True, values other than the default value is not accceptable.


def sticky_attribute_assignment(trackable, name, value):
    if isinstance(value, (tracking.TrackedList, tracking.TrackedDict, tracking.TrackedSet)) and hasattr(trackable, '_tracked'):
        trackable._tracked.append(name)
    if not tracking.is_tracking_enabled():
        return value
    if isinstance(value, tf.__internal__.tracking.Trackable):
        trackable._track_trackable(value, name=name, overwrite=True)
    return value

Adds dependencies, called from __setattr__.

Args:
        trackable: The object to add dependencies to (generally the one having
        an attribute assigned).
        name: The attribute name being assigned.
        value: The value being assigned. Not necessarily a trackable object.

Returns:
        The value which should be stored in the attribute.


def decode(self, spec, encoded_value):
    raise NotImplementedError(f'{type(self).__name__}.decode')

Decodes `value` from a batchable tensor encoding.

Args:
        spec: The TypeSpec for the result value.  If encoded values with spec `s`
        were batched, then `spec` should be `s.batch(batch_size)`; or if encoded
        values with spec `s` were unbatched, then `spec` should be
        `s.unbatch()`.
        encoded_value: A nest of values returned by `encode`; or a nest of values
        that was formed by stacking, unstacking, or concatenating the
        corresponding elements of values returned by `encode`.

Returns:
        A value compatible with `type_spec`.


def json_fhir_string_to_proto(raw_json: str, proto_cls: Type[_T], *, validate: bool=True, default_timezone: str=_primitive_time_utils.SIMPLE_ZULU) -> _T:
    resource = proto_cls()
    merge_json_fhir_string_into_proto(raw_json, resource, validate=validate, default_timezone=default_timezone)
    return resource

Creates a resource of proto_cls and merges contents of raw_json into it.

Args:
        raw_json: The raw FHIR JSON string to convert.
        proto_cls: A subclass of message.Message to instantiate and return.
        validate: A Boolean value indicating if validation should be performed on
        the resultant Message. Validation takes the form of ensuring that basic
        checks such as cardinality guarantees, required field adherence, etc. are
        met. Defaults to True.
        default_timezone: A string specifying the timezone string to use for time-
        like FHIR data during parsing. Defaults to 'Z' for UTC.

Raises:
        fhir_errors.InvalidFhirError: In the event that raw_json was not valid FHIR.

Returns:
        An instance of proto_cls with FHIR JSON data from the raw_json
        representation.


def __init__(self, founded_command='_silence_', score=0.0, is_new_command=False):
    self._founded_command = founded_command
    self._score = score
    self._is_new_command = is_new_command

Construct a recognition result.

Args:
        founded_command: A string indicating the word just founded.
        score: A float representing the confidence of founded word.
        is_new_command: A boolean indicating if the founded command is a new one
        against the last one.


def forward(self, spectrogram: torch.FloatTensor, global_conditioning: Optional[torch.FloatTensor]=None) -> torch.FloatTensor:
    hidden_states = self.conv_pre(spectrogram)
    if global_conditioning is not None:
        hidden_states = hidden_states + self.cond(global_conditioning)
    for i in range(self.num_upsamples):
        hidden_states = nn.functional.leaky_relu(hidden_states, self.config.leaky_relu_slope)
        hidden_states = self.upsampler[i](hidden_states)
        res_state = self.resblocks[i * self.num_kernels](hidden_states)
        for j in range(1, self.num_kernels):
            res_state += self.resblocks[i * self.num_kernels + j](hidden_states)
        hidden_states = res_state / self.num_kernels
    hidden_states = nn.functional.leaky_relu(hidden_states)
    hidden_states = self.conv_post(hidden_states)
    waveform = torch.tanh(hidden_states)
    return waveform

Converts a spectrogram into a speech waveform.

Args:
        spectrogram (`torch.FloatTensor` of shape `(batch_size, config.spectrogram_bins, sequence_length)`):
            Tensor containing the spectrograms.
        global_conditioning (`torch.FloatTensor` of shape `(batch_size, config.speaker_embedding_size, 1)`, *optional*):
            Tensor containing speaker embeddings, for multispeaker models.

Returns:
        `torch.FloatTensor`: Tensor of shape shape `(batch_size, 1, num_frames)` containing the speech waveform.


Called at the beginning of evaluation or validation.

    Subclasses should override for any actions to run.

Args:
        logs: Dict. Currently no data is passed to this argument for this
        method but that may change in the future.


def convert_to_tensor_or_ragged_tensor(value, dtype=None, preferred_dtype=None, name=None):
    if isinstance(value, RaggedTensor):
        if dtype and (not dtype.is_compatible_with(value.dtype)):
            raise ValueError(f'Tensor conversion requested dtype {dtype.name} for RaggedTensor with dtype {value.dtype.name}: {value}.')
        return value
    elif isinstance(value, ragged_tensor_value.RaggedTensorValue):
        with ops.name_scope(name, 'ConvertToTensorOrRaggedTensor', []):
            flat_values = ops.convert_to_tensor(value=value.flat_values, dtype=dtype, dtype_hint=preferred_dtype, name='flat_values')
            return RaggedTensor.from_nested_row_splits(flat_values, value.nested_row_splits, validate=False)
    else:
        return tensor_conversion.convert_to_tensor_v2_with_dispatch(value=value, dtype=dtype, dtype_hint=preferred_dtype, name=name)

Converts value to a `RaggedTensor` or `Tensor`.

    * If `value` is a `RaggedTensor`, then return it as-is.
    * If `value` is a `RaggedTensorValue`, return a corresponding constant
        `RaggedTensor`.
    * Otherwise, use `convert_to_tensor` to convert `value` to a `Tensor`.

Args:
        value: A `RaggedTensor`, a `RaggedTensorValue`, or an object whose type has
        a registered `Tensor` conversion function.
        dtype: Optional element type for the returned tensor.  If missing the type
        is inferred from the type of `value`.
        preferred_dtype: Optional element type for the returned tensor, used when
        dtype is None.  This argument has no effect if `value` is already a
        tensor, or when conversion is not possible.
        name: Optional name to use if a new `Tensor` is created.

Returns:
        A `Tensor` or `RaggedTensor`.


def deconv_output_length(input_length, filter_size, padding, stride):
    if input_length is None:
        return None
    input_length *= stride
    if padding == 'valid':
        input_length += max(filter_size - stride, 0)
    elif padding == 'full':
        input_length -= stride + filter_size - 2
    return input_length

Determines output length of a transposed convolution given input length.

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

Returns:
        The output length (integer).


def forward(self, hidden_states: torch.Tensor, original_hidden_states: Optional[torch.Tensor]=None, layer_idx: Optional[int]=None, attention_mask: Optional[torch.Tensor]=None, causal_mask: Optional[torch.Tensor]=None, past_key_value: Optional[Zamba2HybridDynamicCache]=None, output_attentions: Optional[bool]=False, use_cache: Optional[bool]=False, position_embeddings: Optional[torch.LongTensor]=None) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
    layer_outputs = self.shared_transformer(hidden_states, original_hidden_states=original_hidden_states, layer_idx=layer_idx, attention_mask=causal_mask, past_key_value=past_key_value, output_attentions=output_attentions, position_embeddings=position_embeddings)
    transformer_hidden_states = layer_outputs[0]
    if output_attentions:
        self_attn_weights = layer_outputs[1]
    transformer_hidden_states = self.linear(transformer_hidden_states)
    layer_outputs = self.mamba_decoder(hidden_states, transformer_hidden_states=transformer_hidden_states, attention_mask=attention_mask, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, position_embeddings=position_embeddings)
    if output_attentions:
        layer_outputs = (layer_outputs[0], self_attn_weights) + layer_outputs[2:]
    return layer_outputs

Args:
    hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
    original_hidden_states (`torch.FloatTensor`): word embedding output that will be concatenated with
    hidden activations to form the input of the shared transformer layer.
    layer_idx (`int`): layer number.
    attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
        `(batch, sequence_length)` where padding elements are indicated by 0.
    past_key_value (`Zamba2HybridDynamicCache`, *optional*): cached past key and value projection states
    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.
    use_cache (`bool`, *optional*):
        If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
        (see `past_key_values`).
    position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
        Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
        with `head_dim` being the embedding dimension of each attention head.


def _wrap_and_check_metrics(self, metrics):
    if not isinstance(metrics, dict):
        metrics = {self.METRICS_NAME: metrics}
    outputs = {}
    for key, value in metrics.items():
        if isinstance(value, tuple):
            metric_val, metric_op = value
        else:
            metric_val = value.result()
            assert len(value.updates) == 1
            metric_op = value.updates[0]
        key = self._check_output_key(key, self.METRICS_NAME)
        key = self._prefix_key(key, self.METRICS_NAME)
        val_name = key + self._SEPARATOR_CHAR + self.METRIC_VALUE_SUFFIX
        op_name = key + self._SEPARATOR_CHAR + self.METRIC_UPDATE_SUFFIX
        if not isinstance(metric_val, tensor.Tensor):
            raise ValueError('{} output value must be a Tensor; got {}.'.format(key, metric_val))
        if not (tensor_util.is_tf_type(metric_op) or isinstance(metric_op, ops.Operation)):
            raise ValueError('{} update_op must be a Tensor or Operation; got {}.'.format(key, metric_op))
        metric_op_tensor = metric_op
        if not isinstance(metric_op, tensor.Tensor):
            with ops.control_dependencies([metric_op]):
                metric_op_tensor = constant_op.constant([], name='metric_op_wrapper')
        outputs[val_name] = metric_val
        outputs[op_name] = metric_op_tensor
    return outputs

Handle the saving of metrics.

    Metrics is either a tuple of (value, update_op), or a dict of such tuples.
    Here, we separate out the tuples and create a dict with names to tensors.

Args:
        metrics: Dict of metric results keyed by name.
        The values of the dict can be one of the following:
        (1) instance of `Metric` class.
        (2) (metric_value, update_op) tuples, or a single tuple.
        metric_value must be a Tensor, and update_op must be a Tensor or Op.

Returns:
        dict of output_names to tensors

Raises:
        ValueError: if the dict key is not a string, or the metric values or ops
        are not tensors.


def cos(self: EventSetOrNode) -> EventSetOrNode:
    from temporian.core.operators.unary import cos
    return cos(self)

Calculates the cosine of an [`EventSet`][temporian.EventSet]'s features.

    Can only be used on floating point features.

Example:
        ```python
        >>> a = tp.event_set(
        ...     timestamps=[1, 2, 3, 4, 5],
        ...     features={"M": [0, np.pi/3, np.pi/2, np.pi, 2*np.pi]},
        ... )
        >>> a.cos()
        indexes: ...
                timestamps: [1. 2. 3. 4. 5.]
                'M': [ 1.0000e+00  5.0000e-01  6.1232e-17 -1.0000e+00  1.0000e+00]
        ...

        ```

Returns:
        EventSetOrNode with cosine of input features.


def _list_all_concrete_functions_for_serialization(self):
    seen_signatures = []
    if self.input_signature is not None:
        seen_signatures.append((self.input_signature, {}))
    else:
        concrete_functions = self._list_all_concrete_functions()
        for concrete_function in concrete_functions:
            signature = concrete_function.structured_input_signature
            flattened = nest.flatten(signature)
            if any((isinstance(arg, func_graph_module.UnknownArgument) for arg in flattened)):
                logging.info('Unsupported signature for serialization: %s.', signature)
                continue
            equal_to_signature = functools.partial(function_type_utils.is_same_structure, signature, check_values=True)
            if not any((equal_to_signature(s) for s in seen_signatures)):
                seen_signatures.append(signature)
    concrete_functions = []
    for args, kwargs in seen_signatures:
        concrete_functions.append(self.get_concrete_function(*args, **kwargs))
    return concrete_functions

Returns all concrete functions for serialization.

Returns:
        A list of instances of `ConcreteFunction`.


def execution_to_tensor_values(self, execution):
    debug_event = self._reader.read_execution_event(execution.locator)
    return [_parse_tensor_value(tensor_proto) for tensor_proto in debug_event.execution.tensor_protos]

Read the full tensor values from an Execution or ExecutionDigest.

Args:
        execution: An `ExecutionDigest` or `ExecutionDigest` object.

Returns:
        A list of numpy arrays representing the output tensor values of the
        execution event.


def _client_receive(self):
    try:
        return self._client.readline()
    except socket.error as e:
        raise errors.Error(self._device, f'Encountered socket error "{e}" reading RPC response') from e

Receives the server's response of an RPC message.

Returns:
        Raw bytes of the response.

Raises:
        errors.Error: if a socket error occurred during the read.


def __init__(self, save_steps=None, save_secs=None, output_dir=None, summary_writer=None, scaffold=None, summary_op=None):
    if scaffold is None and summary_op is None or (scaffold is not None and summary_op is not None):
        raise ValueError('Exactly one of scaffold or summary_op must be provided.')
    self._summary_op = summary_op
    self._summary_writer = summary_writer
    self._output_dir = output_dir
    self._scaffold = scaffold
    self._timer = SecondOrStepTimer(every_secs=save_secs, every_steps=save_steps)

Initializes a `SummarySaverHook`.

Args:
        save_steps: `int`, save summaries every N steps. Exactly one of
        `save_secs` and `save_steps` should be set.
        save_secs: `int`, save summaries every N seconds.
        output_dir: `string`, the directory to save the summaries to. Only used if
        no `summary_writer` is supplied.
        summary_writer: `SummaryWriter`. If `None` and an `output_dir` was passed,
        one will be created accordingly.
        scaffold: `Scaffold` to get summary_op if it's not provided.
        summary_op: `Tensor` of type `string` containing the serialized `Summary`
        protocol buffer or a list of `Tensor`. They are most likely an output by
        TF summary methods like `tf.compat.v1.summary.scalar` or
        `tf.compat.v1.summary.merge_all`. It can be passed in as one tensor; if
        more than one, they must be passed in as a list.

Raises:
        ValueError: Exactly one of scaffold or summary_op should be set.


def infer_schema(stats_path, schema_path):
    print('Infering schema from statistics.')
    schema = tfdv.infer_schema(tfdv.load_statistics(stats_path), infer_feature_shape=False)
    print(text_format.MessageToString(schema))
    print('Writing schema to output path.')
    file_io.write_string_to_file(schema_path, text_format.MessageToString(schema))

Infers a schema from stats in stats_path.

Args:
        stats_path: Location of the stats used to infer the schema.
        schema_path: Location where the inferred schema is materialized.


def convert_old_keys_to_new_keys(state_dict_keys: Optional[dict]=None):
    output_dict = {}
    if state_dict_keys is not None:
        old_text = '\n'.join(state_dict_keys)
        new_text = old_text
        for pattern, replacement in ORIGINAL_TO_CONVERTED_KEY_MAPPING.items():
            if replacement is None:
                new_text = re.sub(pattern, '', new_text)
                continue
            new_text = re.sub(pattern, replacement, new_text)
        output_dict = dict(zip(old_text.split('\n'), new_text.split('\n')))
    return output_dict

Converts old keys to new keys using the mapping and dynamically removes the 'ijepa.' prefix if necessary.

Args:
        state_dict_keys (dict): The keys from the state_dict to convert.

Returns:
        dict: A mapping from old keys to new keys.


def should_fire(self, time_domain, timestamp, window, context):
    pass

Whether this trigger should cause the window to fire.

Args:
        time_domain: WATERMARK for event-time timers and REAL_TIME for
        processing-time timers.
        timestamp: for time_domain WATERMARK, it represents the
        watermark: (a lower bound on) the watermark of the system
        and for time_domain REAL_TIME, it represents the
        trigger: timestamp of the processing-time timer.
        window: the window whose trigger is being considered
        context: a context (e.g. a TriggerContext instance) for managing state
        and setting timers

Returns:
        whether this trigger should cause a firing


def Encode(self):
    assert self.value_dict_or_array is not None
    logging.log(1, 'Encoding ' + self.name)
    resolved = MessageValue._ResolveVars(self.value_dict_or_array)
    logging.debug('Resolved: ' + str(resolved))
    return self.msg.encoding.SerializeToString(resolved, self.msg)

Encode this message instance into actual data stream.

    The supported encoding methods are: json, protobuf, and user-defined
    encodings.

Returns:
        A string encoded.


def __call__(self, images: ImageInput=None, text: Optional[Union[str, List[str], TextInput, PreTokenizedInput]]=None, audio=None, videos=None, **kwargs: Unpack[BlipProcessorKwargs]) -> BatchEncoding:
    if images is None and text is None:
        raise ValueError('You have to specify either images or text.')
    text_encoding = None
    output_kwargs = self._merge_kwargs(BlipProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs)
    if text is not None:
        text_encoding = self.tokenizer(text, **output_kwargs['text_kwargs'])
    if images is not None:
        encoding_image_processor = self.image_processor(images, **output_kwargs['images_kwargs'])
        if text_encoding is not None:
            encoding_image_processor.update(text_encoding)
        return encoding_image_processor
    return text_encoding

This method uses [`BlipImageProcessor.__call__`] method to prepare image(s) for the model, and
    [`BertTokenizerFast.__call__`] to prepare text for the model.

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

Args:
        images (`ImageInput`):
            The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
            tensor. Both channels-first and channels-last formats are supported.
        text (`TextInput`, `PreTokenizedInput`, `List[TextInput]`, `List[PreTokenizedInput]`):
            The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
            (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
            `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
        return_tensors (`str` or [`~utils.TensorType`], *optional*):
            If set, will return tensors of a particular framework. Acceptable values are:
                - `'tf'`: Return TensorFlow `tf.constant` objects.
                - `'pt'`: Return PyTorch `torch.Tensor` objects.
                - `'np'`: Return NumPy `np.ndarray` objects.
                - `'jax'`: Return JAX `jnp.ndarray` objects.


def __init__(self, config, exclusions=None):
    self.config = config
    self.stored_ids = {name: set() for name in self.config.keys()}
    self.locked = False
    self._lock_violation_msg = None
    self.exclusions = exclusions or {}

Attribute tracker, used for e.g. Variable tracking.

    Monitors certain attribute types
    and put them in appropriate lists in case of a match.

    Also passively tracks certain mutable collections
    (dict, list) so that items added to them later
    still get tracked. This is done by wrapping these
    collections into an equivalent, tracking-aware object.

Example:
    ```python
    def __init__(self):
        self.tracker = Tracker(
            # Format: `name: (test_fn, store)`
            {
                "variables":
                    (lambda x: isinstance(x, Variable), self._variables),
                "metrics": (lambda x: isinstance(x, Metric), self._metrics),
                "layers": (lambda x: isinstance(x, Layer), self._layers),
            }
        )

    def __setattr__(self, name, value):
        if hasattr(self, "_tracker"):
            value = self._tracker.track(value)
        return super().__setattr__(name, value)
    ```


def _check_params(window_length, dtype):
    if not dtype.is_floating:
        raise ValueError('dtype must be a floating point type. Found %s' % dtype)
    window_length = ops.convert_to_tensor(window_length, dtype=dtypes.int32)
    window_length.shape.assert_has_rank(0)
    return window_length

Check window_length and dtype params.

Args:
        window_length: A scalar value or `Tensor`.
        dtype: The data type to produce. Must be a floating point type.

Returns:
        window_length converted to a tensor of type int32.

Raises:
        ValueError: If `dtype` is not a floating point type or window_length is not
        a scalar.


def __init__(self, unique_id, start_logits, end_logits, start_top_index=None, end_top_index=None, cls_logits=None):
    self.start_logits = start_logits
    self.end_logits = end_logits
    self.unique_id = unique_id
    if start_top_index:
        self.start_top_index = start_top_index
        self.end_top_index = end_top_index
        self.cls_logits = cls_logits

Constructs a SquadResult which can be used to evaluate a model's output on the SQuAD dataset.

Args:
        unique_id: The unique identifier corresponding to that example.
        start_logits: The logits corresponding to the start of the answer
        end_logits: The logits corresponding to the end of the answer


def call_with_mapped_args(self, mapped_args: MappedArgs[FrameType]) -> _HasReturnT:

Calls this function with the given mapped arguments.

Args:
        mapped_args: The function arguments mapped to parameter names.

Returns:
        An object with information about the result of the function call, with a
        get_return_value() method that retrieves the return value.


def _pad_for_batching(self, pixel_values: List[np.ndarray], data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None):
    max_patch = max((len(x) for x in pixel_values))
    pixel_values = [self.pad(image, padding=((0, max_patch - image.shape[0]), (0, 0), (0, 0), (0, 0)), data_format=data_format, input_data_format=input_data_format) for image in pixel_values]
    return pixel_values

Pads images on the `num_of_patches` dimension with zeros to form a batch of same number of patches.

Args:
        pixel_values (`List[np.ndarray]`):
            An array of pixel values of each images of shape (`batch_size`, `num_patches`, `image_in_3D`)
        data_format (`str` or `ChannelDimension`, *optional*):
            The channel dimension format for the output image. Can be one of:
                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
            If unset, will use same as the input image.
        input_data_format (`str` or `ChannelDimension`, *optional*):
            The channel dimension format for the input image. Can be one of:
                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
            If unset, will use the inferred format of the input image.

Returns:
        List[`np.ndarray`]: The padded images.


def validate_slicing_string(slicing_string):
    return bool(re.search('^\\[(\\d|,|\\s|:)+\\]$', slicing_string))

Validate a slicing string.

    Check if the input string contains only brackets, digits, commas and
    colons that are valid characters in numpy-style array slicing.

Args:
        slicing_string: (str) Input slicing string to be validated.

Returns:
        (bool) True if and only if the slicing string is valid.


def volatility(self, strike: types.FloatTensor, expiry_dates: Optional[types.DateTensor]=None, expiry_times: Optional[types.FloatTensor]=None, term: Optional[types.Period]=None) -> types.FloatTensor:
    pass

Returns the interpolated volatility on a specified set of expiries.

Args:
        strike: The strikes for which the interpolation is desired.
        expiry_dates: Optional input specifying the expiry dates for which
        interpolation is desired. The user should supply either `expiry_dates`
        or `expiry_times` for interpolation.
        expiry_times: Optional real `Tensor` containing the time to expiration
        for which interpolation is desired. The user should supply either
        `expiry_dates` or `expiry_times` for interpolation.
        term: Optional input specifying the term of the underlying rate for
        which the interpolation is desired. Relevant for interest rate implied
        volatility data.

Returns:
        A `Tensor` of the same shape as `expiry` with the interpolated volatility
        from the volatility surface.


def _BiasAddGradV1(unused_bias_op: ops.Operation, received_grad):
    reduction_dim_tensor = math_ops.range(array_ops.rank(received_grad) - 1)
    return (received_grad, math_ops.reduce_sum(received_grad, reduction_dim_tensor))

Return the gradients for the 2 inputs of bias_op.

    The first input of unused_bias_op is the tensor t, and its gradient is
    just the gradient the unused_bias_op received.

    The second input of unused_bias_op is the bias vector which has one fewer
    dimension than "received_grad" (the batch dimension.)  Its gradient is the
    received gradient Summed on the batch dimension, which is the first dimension.

Args:
        unused_bias_op: The BiasOp for which we need to generate gradients.
        received_grad: Tensor.  The gradients passed to the BiasOp.

Returns:
        Two tensors, the first one for the "tensor" input of the BiasOp,
        the second one for the "bias" input of the BiasOp.


def update_state(self, y_true, y_pred, sample_weight=None):
    metrics_utils.update_confusion_matrix_variables({metrics_utils.ConfusionMatrix.TRUE_POSITIVES: self.true_positives, metrics_utils.ConfusionMatrix.FALSE_NEGATIVES: self.false_negatives}, y_true, y_pred, thresholds=self.thresholds, thresholds_distributed_evenly=self._thresholds_distributed_evenly, top_k=self.top_k, class_id=self.class_id, sample_weight=sample_weight)

Accumulates true positive and false negative statistics.

Args:
        y_true: The ground truth values, with the same dimensions as
            `y_pred`. Will be cast to `bool`.
        y_pred: The predicted values. Each element must be in the range
            `[0, 1]`.
        sample_weight: Optional weighting of each example. Defaults to `1`.
            Can be a tensor whose rank is either 0, or the same rank as
            `y_true`, and must be broadcastable to `y_true`.


def forward(self, input_tensor: torch.LongTensor, attention_mask: Optional[bool]=None, output_hidden_states: Optional[bool]=None, output_attentions: Optional[bool]=False, return_dict: Optional[bool]=None):
    feature_representation = input_tensor
    if self.embed is not None:
        feature_representation = self.embed(feature_representation)
    hidden_states, pos_emb = self.pos_enc(feature_representation)
    all_hidden_states = () if output_hidden_states else None
    all_self_attentions = () if output_attentions else None
    for conformer_layer in self.conformer_layers:
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)
        layer_outputs = conformer_layer(hidden_states, pos_emb, attention_mask, output_attentions)
        hidden_states = layer_outputs[0]
        if output_attentions:
            all_self_attentions = all_self_attentions + (layer_outputs[1],)
    if output_hidden_states:
        all_hidden_states = all_hidden_states + (hidden_states,)
    if not return_dict:
        return tuple((v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None))
    return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_self_attentions)

Args:
    input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
        Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
        provide it.

        Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
        [`PreTrainedTokenizer.__call__`] for details.

        [What are input IDs?](../glossary#input-ids)
    attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
        Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

        - 1 for tokens that are **not masked**,
        - 0 for tokens that are **masked**.

        [What are attention masks?](../glossary#attention-mask)
    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.
    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.
    return_dict (`bool`, *optional*):
        Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.

Returns:
    `torch.Tensor`:
        Output tensor of shape `(batch, time, attention_dim)`.


def _get_executor_init(self, workers):

    def pool_fn(seqs):
        pool = get_pool_class(True)(workers, initializer=init_pool_generator, initargs=(seqs, self.random_seed, get_worker_id_queue()))
        _DATA_POOLS.add(pool)
        return pool
    return pool_fn

Gets the Pool initializer for multiprocessing.

Args:
        workers: Number of works.

Returns:
        A Function to initialize the pool


def add_ordinary_node(self, ast_node):
    node = self._add_new_node(ast_node)
    self.leaves = set((node,))
    return node

Grows the graph by adding an ordinary CFG node.

    Ordinary nodes are followed by the next node, in lexical order, that is,
    they become the new leaf set.

Args:
        ast_node: ast.AST

Returns:
        Node


def _preprocess_grad(grad, body_graph_output, while_op_input, while_op_output):
    if not _is_trainable(body_graph_output):
        return None
    if while_op_output.dtype in (dtypes.resource, dtypes.variant) and default_gradient.supports_default_grad(while_op_input) and (grad is None):
        return _zeros_like(while_op_input, while_op_output)
    if isinstance(grad, indexed_slices.IndexedSlices):
        return ops.convert_to_tensor(grad)
    return grad

Returns the initial gradient to be used for a given output tensor.

Args:
        grad: the original gradient Tensor passed to the gradient function.
        body_graph_output: the corresponding Tensor in the body graph.
        while_op_input: the corresponding Tensor input of the While op.
        while_op_output: the corresponding Tensor output of the While op.

Returns:
        A Tensor or None.


def find_all_documented_objects() -> List[str]:
    documented_obj = []
    documented_methods_map = {}
    for doc_file in Path(PATH_TO_DOC).glob('**/*.md'):
        with open(doc_file, 'r', encoding='utf-8', newline='\n') as f:
            content = f.read()
        raw_doc_objs = re.findall('\\[\\[autodoc\\]\\]\\s+(\\S+)\\s+', content)
        documented_obj += [obj.split('.')[-1] for obj in raw_doc_objs]
        for obj in raw_doc_objs:
            obj_public_methods = re.findall(f'\\[\\[autodoc\\]\\] {obj}((\\n\\s+-.*)+)', content)
            if len(obj_public_methods) == 0:
                continue
            else:
                documented_methods_map[obj] = re.findall('(?<=-\\s).*', obj_public_methods[0][0])
    return (documented_obj, documented_methods_map)

Parse the content of all doc files to detect which classes and functions it documents.

Returns:
        `List[str]`: The list of all object names being documented.
        `Dict[str, List[str]]`: A dictionary mapping the object name (full import path, e.g.
            `integrations.PeftAdapterMixin`) to its documented methods


def __init__(self, variables, name='ShardedVariable'):
    super(ShardedVariableMixin, self).__init__()
    self._variables = variables
    self._name = name
    if not isinstance(variables, Sequence) or not variables or any((not isinstance(v, variables_lib.Variable) for v in variables)):
        raise TypeError(f'Argument `variables` should be a non-empty list of `variables.Variable`s. Received {variables}')
    var_dtypes = {v.dtype for v in variables}
    if len(var_dtypes) > 1:
        raise ValueError(f'All elements in argument `variables` must have the same dtype. Received dtypes: {[v.dtype for v in variables]}')
    first_var = variables[0]
    self._dtype = first_var.dtype
    higher_dim_shapes = {tuple(v.shape.as_list()[1:]) for v in variables}
    if len(higher_dim_shapes) > 1:
        raise ValueError(f'All elements in argument `variables` must have the same shapes except for the first axis. Received shapes: {[v.shape for v in variables]}')
    first_dim = sum((int(v.shape.as_list()[0]) for v in variables))
    self._shape = tensor_shape.TensorShape([first_dim] + first_var.shape.as_list()[1:])
    for v in variables:
        v._sharded_container = weakref.ref(self)
    self._var_offsets = [[0 for _ in range(len(first_var.shape))] for _ in range(len(variables))]
    for i in range(1, len(variables)):
        self._var_offsets[i][0] += self._var_offsets[i - 1][0] + variables[i - 1].shape.as_list()[0]
    save_slice_info = [v._get_save_slice_info() for v in variables]
    if any((slice_info is not None for slice_info in save_slice_info)):
        raise ValueError(f'`SaveSliceInfo` should not be set for all elements in argument `variables`. `ShardedVariable` will infer `SaveSliceInfo` according to the order of the elements `variables`. Received save slice info {save_slice_info}')
    self._saving_variable = resource_variable_ops.UninitializedVariable(shape=self._shape, dtype=self._dtype, name=self._name, trainable=self._variables[0].trainable, synchronization=variables_lib.VariableSynchronization.NONE, aggregation=variables_lib.VariableAggregation.NONE)

Treats `variables` as shards of a larger Variable.

Example:
    ```
    variables = [
        tf.Variable(..., shape=(10, 100), dtype=tf.float32),
        tf.Variable(..., shape=(15, 100), dtype=tf.float32),
        tf.Variable(..., shape=(5, 100), dtype=tf.float32)
    ]
    sharded_variable = ShardedVariableMixin(variables)
    assert sharded_variable.shape.as_list() == [30, 100]
    ```

Args:
        variables: A list of `ResourceVariable`s that comprise this sharded
        variable. Variables should not be shared between different
        `ShardedVariableMixin` objects.
        name: String. Name of this container. Defaults to "ShardedVariable".


def read_up_to(self, queue, num_records, name=None):
    if isinstance(queue, tensor_lib.Tensor):
        queue_ref = queue
    else:
        queue_ref = queue.queue_ref
    if self._reader_ref.dtype == dtypes.resource:
        return gen_io_ops.reader_read_up_to_v2(self._reader_ref, queue_ref, num_records, name=name)
    else:
        old_queue_op = gen_data_flow_ops.fake_queue(queue_ref)
        return gen_io_ops.reader_read_up_to(self._reader_ref, old_queue_op, num_records, name=name)

Returns up to num_records (key, value) pairs produced by a reader.

    Will dequeue a work unit from queue if necessary (e.g., when the
    Reader needs to start reading from a new file since it has
    finished with the previous file).
    It may return less than num_records even before the last batch.

Args:
        queue: A Queue or a mutable string Tensor representing a handle
        to a Queue, with string work items.
        num_records: Number of records to read.
        name: A name for the operation (optional).

Returns:
        A tuple of Tensors (keys, values).
        keys: A 1-D string Tensor.
        values: A 1-D string Tensor.


def check_alive(self, worker_name):
    if self._context_handle:
        return pywrap_tfe.TFE_ContextCheckAlive(self._context_handle, worker_name)
    else:
        raise ValueError('Context is not initialized.')

Checks whether a remote worker is alive or not.

Args:
        worker_name: a string representing the remote worker. It must be a fully
        specified name like "/job:worker/replica:0/task:0".

Returns:
        a boolean indicating whether the remote worker is alive or not.

Raises:
        ValueError: if context is not initialized.


def upload_file(filename: str, config: Config, full_table_id: str, action: Action=Action.APPEND, service_account_email: Optional[str]=None) -> None:
    if service_account_email:
        auth_config = AuthConfig(service_account_email=service_account_email)
    else:
        auth_config = None
    table_metadata = build_table_metadata(full_table_id)
    credentials = get_credentials(auth_config)
    bq_legacy_client = get_bq_legacy_client(table_metadata.project_id, credentials)
    table_exists = get_table(bq_legacy_client, table_metadata) is not None
    if table_exists and action == Action.REPLACE:
        bq_legacy_client.delete_table(table_metadata.full_table_id)
        table_exists = False
    if not table_exists:
        schema = generate_bigquery_schema(config)
        create_table(bq_legacy_client, table_metadata, schema)
    row_count = get_csv_row_count(filename)
    columns = [column.bq_name for column in config]
    with open(filename, 'r') as f:
        csv = DictReader(f, fieldnames=columns)
        buffer_: List[Row] = []
        buffer_size = row_count // 100
        buffer = Buffer[Row](buffer_, buffer_size, lambda rows: upload_rows(bq_legacy_client, table_metadata, rows))
        for row in csv:
            buffer.push(row)
        buffer.flush(force=True)

Upload a data file conforming to the given config to BigQuery.

Args:
        * filename: Local path to csv file to be uploaded
        * config: Config key from configs.CONFIGS dictionary
        * full_table_id: BigQuery table id
        * action: APPEND to table or REPLACE table
        * service_account_email: Email address of service account


def _compute_llama3_parameters(config: PretrainedConfig, device: 'torch.device', seq_len: Optional[int]=None, **rope_kwargs) -> tuple['torch.Tensor', float]:
    inv_freq, attention_factor = _compute_default_rope_parameters(config, device, seq_len, **rope_kwargs)
    factor = config.rope_scaling['factor']
    low_freq_factor = config.rope_scaling['low_freq_factor']
    high_freq_factor = config.rope_scaling['high_freq_factor']
    old_context_len = config.rope_scaling['original_max_position_embeddings']
    low_freq_wavelen = old_context_len / low_freq_factor
    high_freq_wavelen = old_context_len / high_freq_factor
    wavelen = 2 * math.pi / inv_freq
    inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor, inv_freq)
    smooth_factor = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
    smoothed_inv_freq = (1 - smooth_factor) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
    is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen)
    inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama)
    return (inv_freq_llama, attention_factor)

Computes the inverse frequencies for llama 3.1.

Args:
        config ([`~transformers.PretrainedConfig`]):
            The model configuration.
        device (`torch.device`):
            The device to use for initialization of the inverse frequencies.
        seq_len (`int`, *optional*):
            The current sequence length. Unused for this type of RoPE.
        rope_kwargs (`Dict`, *optional*):
            BC compatibility with the previous RoPE class instantiation, will be removed in v4.45.

Returns:
        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
        post-processing scaling factor applied to the computed cos/sin.


def sparsemax(x, axis=-1):
    if any_symbolic_tensors((x,)):
        return Sparsemax(axis).symbolic_call(x)
    return backend.nn.sparsemax(x, axis=axis)

Sparsemax activation function.

    For each batch `i`, and class `j`,
    sparsemax activation function is defined as:

    `sparsemax(x)[i, j] = max(x[i, j] - τ(x[i, :]), 0).`

Args:
        x: Input tensor.
        axis: `int`, axis along which the sparsemax operation is applied.

Returns:
        A tensor, output of sparsemax transformation. Has the same type and
        shape as `x`.

Example:
    >>> x = np.array([-1., 0., 1.])
    >>> x_sparsemax = keras.ops.sparsemax(x)
    >>> print(x_sparsemax)
    array([0., 0., 1.], shape=(3,), dtype=float64)


def make_collective(self, num_processes, gpu_per_process):
    cluster_resolver = cluster_resolver_lib.TFConfigClusterResolver()
    devices = ['/job:worker/replica:0/task:%d/device:CPU:0' % cluster_resolver.task_id]
    if gpu_per_process > 0:
        devices = ['/job:worker/replica:0/task:%d/device:GPU:%d' % (cluster_resolver.task_id, i) for i in range(gpu_per_process)]
    group_size = num_processes * len(devices)
    collective = cross_device_ops_lib.CollectiveAllReduce(devices=devices, group_size=group_size, options=collective_util.Options())
    return (collective, devices, cluster_resolver.task_id)

Returns collectives and other info to be used in tests.

Args:
        num_processes: an integer indicating the number of processes that
        participate in the collective.
        gpu_per_process: number of GPUs (0 if no GPUs) used by each process.

Returns:
    A tuple of (collective, devices, pid) where collective is a instance
    of `CollectiveAllReduce`, devices are a list of local devices (str)
    attached to the current process, and pid is the id of this process among
    all participant processes.


def _convert(value, dtype=None):
    result = numpy_compat.np_asarray(value, dtype=dtype, order='C')
    if result.dtype.char == 'S' and result is not value:
        return numpy_compat.np_asarray(value, order='C', dtype=object)
    elif result.dtype.char == 'U' and result is not value:
        value = np.vectorize(lambda x: x.encode('utf8'))(value)
        return numpy_compat.np_asarray(value, order='C', dtype=object)
    elif result.dtype.char == 'U':
        return result.astype(np.bytes_)
    else:
        return result

Converts an arg to numpy, avoiding dangerous string and unicode dtypes.

    Numpy pads with zeros when using string and unicode dtypes if different
    components of a tensor have different lengths.  This is bad: ignoring the
    padding is wrong for text data, and removing the padding is wrong for binary
    data.  To avoid this bug, we redo the conversion using an object dtype.
    Additionally, we convert unicode strings to (byte-)strings for
    compatibility.

Args:
        value: Value to convert to a numpy array.
        dtype: (Optional.) Desired NumPy type for the returned value.

Returns:
        A numpy array.


def insert_model_in_doc_toc(old_model_patterns, new_model_patterns):
    toc_file = REPO_PATH / 'docs' / 'source' / 'en' / '_toctree.yml'
    with open(toc_file, 'r', encoding='utf8') as f:
        content = yaml.safe_load(f)
    api_idx = 0
    while content[api_idx]['title'] != 'API':
        api_idx += 1
    api_doc = content[api_idx]['sections']
    model_idx = 0
    while api_doc[model_idx]['title'] != 'Models':
        model_idx += 1
    model_doc = api_doc[model_idx]['sections']
    old_model_type = old_model_patterns.model_type
    section_idx = 0
    while section_idx < len(model_doc):
        sections = [entry['local'] for entry in model_doc[section_idx]['sections']]
        if f'model_doc/{old_model_type}' in sections:
            break
        section_idx += 1
    if section_idx == len(model_doc):
        old_model = old_model_patterns.model_name
        new_model = new_model_patterns.model_name
        print(f'Did not find {old_model} in the table of content, so you will need to add {new_model} manually.')
        return
    toc_entry = {'local': f'model_doc/{new_model_patterns.model_type}', 'title': new_model_patterns.model_name}
    model_doc[section_idx]['sections'].append(toc_entry)
    model_doc[section_idx]['sections'] = sorted(model_doc[section_idx]['sections'], key=lambda s: s['title'].lower())
    api_doc[model_idx]['sections'] = model_doc
    content[api_idx]['sections'] = api_doc
    with open(toc_file, 'w', encoding='utf-8') as f:
        f.write(yaml.dump(content, allow_unicode=True))

Insert the new model in the doc TOC, in the same section as the old model.

Args:
        old_model_patterns (`ModelPatterns`): The patterns for the old model.
        new_model_patterns (`ModelPatterns`): The patterns for the new model.


def __init__(self, eos_token_id: Optional[torch.Tensor]=None, num_output_tokens: int=10, max_matching_ngram_size: Optional[int]=None, max_length: int=20):
    self.num_output_tokens = num_output_tokens
    self.max_matching_ngram_size = max_matching_ngram_size if max_matching_ngram_size else 2
    self.max_length = max_length
    self.eos_token_id = eos_token_id
    if self.max_matching_ngram_size <= 0 or self.num_output_tokens <= 0:
        raise ValueError('Invalid max_matching_ngram_size or num_output_tokens')

`CandidateGenerator` class to be used for prompt lookup generation. This class generates candidates by looking up
    likely continuations in the provided prompt (input_ids) itself.
    Read the following blog post for more information: https://github.com/apoorvumang/prompt-lookup-decoding

Args:
        max_matching_ngram_size (`int`):
            The maximum ngram size to be considered for matching in the prompt
        num_output_tokens (`int`):
            The number of tokens to be output as candidate tokens.
        max_length (`int`):
            The number of total maximum tokens that can be generated. For decoder-only models that includes the prompt length.
            Defaults to 20, which is the max length used as default in generation config.


def ask_to_proceed_with_overwrite(filepath):
    overwrite = input('[WARNING] %s already exists - overwrite? [y/n]' % filepath).strip().lower()
    while overwrite not in ('y', 'n'):
        overwrite = input('Enter "y" (overwrite) or "n" (cancel).').strip().lower()
    if overwrite == 'n':
        return False
    print('[TIP] Next time specify overwrite=True!')
    return True

Produces a prompt asking about overwriting a file.

Args:
        filepath: the path to the file to be overwritten.

Returns:
        True if we can proceed with overwrite, False otherwise.


def _to_backend_layout(tensor_layout):
    if tensor_layout.device_mesh is None:
        raise ValueError('Cannot create sharding when device mesh is not set for TensorLayout.')
    sharding_specs = [axis if axis else dtensor.UNSHARDED for axis in tensor_layout.axes]
    dtensor_mesh = tensor_layout.device_mesh.backend_mesh
    return dtensor.Layout(sharding_specs=sharding_specs, mesh=dtensor_mesh)

Convert the TensorLayout to Tensorflow backend specific Sharding.

Args:
        tensor_layout: TensorLayout instance to convert.

Returns:
        A `tf.dtensor.Layout` instance.


def get_input_mask_at(self, node_index):
    inputs = self.get_input_at(node_index)
    if isinstance(inputs, list):
        return [getattr(x, '_keras_mask', None) for x in inputs]
    else:
        return getattr(inputs, '_keras_mask', None)

Retrieves the input mask tensor(s) of a layer at a given node.

Args:
        node_index: Integer, index of the node
            from which to retrieve the attribute.
            E.g. `node_index=0` will correspond to the
            first time the layer was called.

Returns:
        A mask tensor
        (or list of tensors if the layer has multiple inputs).


def prune(self, feeds, fetches, name=None, input_signature=None, are_keyword_args_also_positional=False):
    name = name or 'pruned'
    flat_feeds = nest.flatten(feeds, expand_composites=True)
    flat_feeds = [self.graph.as_graph_element(t) for t in flat_feeds]
    for f in flat_feeds:
        if not isinstance(f, tensor_lib.Tensor):
            raise ValueError(f'All members of argument `feeds` must be tensors. Got {f} with type {type(f)}.')
    internal_captures = {id(c) for c in self.graph.internal_captures}
    flat_feeds = [f for f in flat_feeds if id(f) not in internal_captures]
    operation_fetches = []
    tensor_fetches = []
    tensor_infos = []

    def _fetch_preprocessing_callback(fetch):
        """Extract out lists of ops, tensors, and tensor type info.

      Turns TensorInfos into Tensors in the original `fetches` structure.
      Also extracts ops from `fetches`.

      Args:
        fetch: The fetch to preprocess: Tensor, TensorInfo, or Operation, or
          string identifying a Tensor or Operation.

      Returns:
        `fetch` converted to a Tensor.
      """
        if isinstance(fetch, ops.Operation):
            operation_fetches.append(fetch)
            return fetch
        elif isinstance(fetch, meta_graph_pb2.TensorInfo):
            tensor_infos.append(fetch)
            decoded = _get_element_from_tensor_info(fetch, self._func_graph)
            if tensor_util.is_tf_type(decoded) or isinstance(decoded, composite_tensor.CompositeTensor):
                tensor_fetches.append(decoded)
            else:
                operation_fetches.append(decoded)
            return decoded
        elif isinstance(fetch, (tensor_lib.Tensor, composite_tensor.CompositeTensor)):
            tensor_fetches.append(fetch)
            return fetch
        else:
            graph_element = self.graph.as_graph_element(fetch)
            return _fetch_preprocessing_callback(graph_element)
    fetches = nest.map_structure(_fetch_preprocessing_callback, fetches)
    tensor_fetches = nest.flatten(tensor_fetches, expand_composites=True)
    for f in flat_feeds + tensor_fetches + operation_fetches:
        if f.graph is not self._func_graph:
            raise ValueError(f'Can only prune function whose feeds and fetches from graph {self._func_graph}. Input {f} is from a different graph {f.graph}.')
    with self._func_graph.as_default():
        pruned_graph = func_graph.FuncGraph(name)
    lift_map = lift_to_graph.lift_to_graph(operation_fetches + tensor_fetches, pruned_graph, sources=flat_feeds + self.graph.internal_captures, base_graph=self._func_graph)
    pruned_graph.outputs.extend((lift_map[x] for x in tensor_fetches))
    pruned_graph.control_outputs.extend([lift_map[operation] for operation in operation_fetches])
    pruned_graph.inputs.extend((lift_map[x] for x in flat_feeds))
    for external_capture, internal_capture in self.graph.captures:
        pruned_graph.add_capture(external_capture, lift_map[internal_capture])
    for ti in tensor_infos:
        if ti.WhichOneof('encoding') == 'name':
            t = pruned_graph.as_graph_element(ti.name)
            if tensor_util.is_tf_type(t):
                t.set_shape(tensor_shape.TensorShape(ti.tensor_shape))
    for f in self.graph._functions.values():
        pruned_graph._add_function(f)
    pruned_graph.variables = self.graph.variables

    def _structured_output_mapping(fetched):
        """callback for `nest.map_structure()`"""
        lifted = lift_map[fetched]
        if isinstance(lifted, ops.Operation):
            return None
        return lifted
    pruned_graph.structured_outputs = nest.map_structure(_structured_output_mapping, fetches, expand_composites=True)
    if input_signature:
        args, kwargs = input_signature
        args = () if args is None else args
        input_signature = (args, kwargs)
    pruned_graph.structured_input_signature = input_signature
    pruned_fn = WrappedFunction(pruned_graph, variable_holder=self._variable_holder, are_keyword_args_also_positional=are_keyword_args_also_positional)
    pruned_fn._num_positional_args = len(flat_feeds)
    pruned_fn._arg_keywords = [tensor.op.name for tensor in flat_feeds]
    return pruned_fn

Extract a subgraph of this function's underlying graph.

    Wraps the subgraph in a new `WrappedFunction` object.

Args:
        feeds: Input tensors to the subgraph to extract, as `Tensor` objects.
        fetches: Possibly-nested Python data structure containing information
        about outputs of the target subgraph. Each entry can either be a
        `Tensor` object (for data outputs), an `Operation` object (for control
        outputs), or a `TensorInfo` proto. Any additional shape/dtype
        information provided in a `TensorInfo` and not present in the original
        graph will be added to the returned subgraph.
        name: (optional) Name to give to the underlying `FuncGraph` of the
        returned object. If no name is provided, the graph's name will be
        `"pruned"`.
        input_signature: (optional) possibly-nested Python data structure
        containing `TensorSpec` objects, with which to populate the returned
        functions's `FuncGraph`'s `structured_input_signature` field.
        are_keyword_args_also_positional: whether the keyword arguments
        in `input_signature` are `POSITIONAL_OR_KEYWORD` arguments. If
        `False` (default), they are treated as `KEYWORD_ONLY`
        arguments.

Returns:
        A new `WrappedFunction` object containing a copy of the portion of this
        object's graph that goes from `feeds` to `fetches`.


def _get_node_attribute_at_index(self, node_index, attr, attr_name):
    if not self._inbound_nodes:
        raise RuntimeError('The layer has never been called and thus has no defined ' + attr_name + '.')
    if not len(self._inbound_nodes) > node_index:
        raise ValueError('Asked to get ' + attr_name + ' at node ' + str(node_index) + ', but the layer has only ' + str(len(self._inbound_nodes)) + ' inbound nodes.')
    values = getattr(self._inbound_nodes[node_index], attr)
    if isinstance(values, list) and len(values) == 1:
        return values[0]
    else:
        return values

Private utility to retrieves an attribute (e.g. inputs) from a node.

    This is used to implement the methods:
        - get_input_shape_at
        - get_output_shape_at
        - get_input_at
        etc...

Args:
        node_index: Integer index of the node from which
            to retrieve the attribute.
        attr: Exact node attribute name.
        attr_name: Human-readable attribute name, for error messages.

Returns:
        The layer's attribute `attr` at the node of index `node_index`.

Raises:
        RuntimeError: If the layer has no inbound nodes, or if called in Eager
        mode.
        ValueError: If the index provided does not match any node.


def call(self, input_ids: tf.Tensor, pixel_values: tf.Tensor | None=None, decoder_input_ids: tf.Tensor | None=None, decoder_attention_mask: tf.Tensor | None=None, attention_mask: tf.Tensor | None=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, labels: tf.Tensor | None=None, return_dict: Optional[bool]=None, training: Optional[bool]=None) -> Union[Tuple, TFBlipTextVisionModelOutput]:
    if labels is None and decoder_input_ids is None:
        raise ValueError('Either `decoder_input_ids` or `labels` should be passed when calling `TFBlipForQuestionAnswering`. if you are training the model make sure that `labels` is passed, if you are using the model for inference make sure that `decoder_input_ids` is passed or call `generate`')
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    vision_outputs = self.vision_model(pixel_values=pixel_values, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, training=training)
    image_embeds = vision_outputs[0]
    image_attention_mask = tf.ones(shape_list(image_embeds)[:-1], dtype=tf.int64)
    question_embeds = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask, encoder_hidden_states=image_embeds, encoder_attention_mask=image_attention_mask, return_dict=return_dict, training=training)
    question_embeds = question_embeds[0] if not return_dict else question_embeds.last_hidden_state
    if labels is not None and decoder_input_ids is None:
        decoder_input_ids = labels
    answer_output = self.text_decoder(input_ids=decoder_input_ids, attention_mask=decoder_attention_mask, encoder_hidden_states=question_embeds, encoder_attention_mask=attention_mask, labels=labels, return_dict=return_dict, training=training)
    if labels is not None:
        decoder_loss = tf.reduce_mean(answer_output.loss) if return_dict else tf.reduce_mean(answer_output[0])
    else:
        decoder_loss = None
    if not return_dict:
        outputs = (decoder_loss, image_embeds, vision_outputs[0]) + vision_outputs[2:]
        return tuple((output for output in outputs if output is not None))
    return TFBlipTextVisionModelOutput(loss=decoder_loss, image_embeds=image_embeds, last_hidden_state=vision_outputs.last_hidden_state, hidden_states=vision_outputs.hidden_states, attentions=vision_outputs.attentions)

Returns:

Example:
    ```python
    >>> from PIL import Image
    >>> import requests
    >>> from transformers import AutoProcessor, TFBlipForQuestionAnswering

    >>> model = TFBlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base")
    >>> processor = AutoProcessor.from_pretrained("Salesforce/blip-vqa-base")

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

    >>> # training
    >>> text = "How many cats are in the picture?"
    >>> label = "2"
    >>> inputs = processor(images=image, text=text, return_tensors="tf")
    >>> labels = processor(text=label, return_tensors="tf").input_ids

    >>> inputs["labels"] = labels
    >>> outputs = model(**inputs)
    >>> loss = outputs.loss

    >>> # inference
    >>> text = "How many cats are in the picture?"
    >>> inputs = processor(images=image, text=text, return_tensors="tf")
    >>> outputs = model.generate(**inputs)
    >>> print(processor.decode(outputs[0], skip_special_tokens=True))
    2
    ```


def __init__(self, pool_size, strides, padding='valid', data_format='channels_last', name=None, **kwargs):
    if strides is None:
        raise ValueError('Argument `strides` must not be None.')
    super(MaxPooling3D, self).__init__(pool_size=pool_size, strides=strides, padding=padding, data_format=data_format, name=name, **kwargs)

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

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


def _pyval_update_fields(pyval, fields, depth):
    if not isinstance(pyval, (dict, list, tuple)):
        raise ValueError('Expected dict or nested list/tuple of dict')
    for key, target in fields.items():
        for _ in range(1, depth):
            target = target[-1]
        target.append(pyval[key] if isinstance(pyval, dict) else [])
    if isinstance(pyval, (list, tuple)):
        for child in pyval:
            _pyval_update_fields(child, fields, depth + 1)

Append the field values from `pyval` to `fields`.

Args:
        pyval: A python `dict`, or nested list/tuple of `dict`, whose value(s)
        should be appended to `fields`.
        fields: A dictionary mapping string keys to field values.  Field values
        extracted from `pyval` are appended to this dictionary's values.
        depth: The depth at which `pyval` should be appended to the field values.


def _relocate_if_symbolic(self, key: Union[str, int], value: Any) -> Any:
    if isinstance(value, Symbolic):
        root_path = utils.KeyPath(key, self.sym_path)
        if value.sym_parent is not None and (value.sym_parent is not self or root_path != value.sym_path):
            value = value.clone()
    if isinstance(value, TopologyAware):
        value.sym_setpath(utils.KeyPath(key, self.sym_path))
        value.sym_setparent(self._sym_parent_for_children())
    return value

Relocate if a symbolic value is to be inserted as member.

    NOTE(daiyip): when a symbolic value is inserted into the object tree,
    if it already has a parent, we need to make a shallow copy of this object
    to avoid multiple parents. Otherwise we need to set its parent and root_path
    according to current object.

Args:
        key: Key used to insert the value.
        value: formal value to be inserted.

Returns:
        Formalized value that is ready for insertion as members.


def plot_model(model, to_file='model.png', show_shapes=False, show_dtype=False, show_layer_names=False, rankdir='TB', expand_nested=False, dpi=200, show_layer_activations=False, show_trainable=False, **kwargs):
    if not model.built:
        raise ValueError('This model has not yet been built. Build the model first by calling `build()` or by calling the model on a batch of data.')
    if not check_pydot():
        message = 'You must install pydot (`pip install pydot`) for `plot_model` to work.'
        if 'IPython.core.magics.namespace' in sys.modules:
            io_utils.print_msg(message)
            return
        else:
            raise ImportError(message)
    if not check_graphviz():
        message = 'You must install graphviz (see instructions at https://graphviz.gitlab.io/download/) for `plot_model` to work.'
        if 'IPython.core.magics.namespace' in sys.modules:
            io_utils.print_msg(message)
            return
        else:
            raise ImportError(message)
    if kwargs.pop('layer_range', None) is not None:
        raise ValueError('Argument `layer_range` is no longer supported.')
    if kwargs:
        raise ValueError(f'Unrecognized keyword arguments: {kwargs}')
    dot = model_to_dot(model, show_shapes=show_shapes, show_dtype=show_dtype, show_layer_names=show_layer_names, rankdir=rankdir, expand_nested=expand_nested, dpi=dpi, show_layer_activations=show_layer_activations, show_trainable=show_trainable)
    to_file = str(to_file)
    if dot is None:
        return
    _, extension = os.path.splitext(to_file)
    if not extension:
        extension = 'png'
    else:
        extension = extension[1:]
    dot.write(to_file, format=extension)
    if extension != 'pdf':
        try:
            from IPython import display
            return display.Image(filename=to_file)
        except ImportError:
            pass

Converts a Keras model to dot format and save to a file.

Example:
    ```python
    inputs = ...
    outputs = ...
    model = keras.Model(inputs=inputs, outputs=outputs)

    dot_img_file = '/tmp/model_1.png'
    keras.utils.plot_model(model, to_file=dot_img_file, show_shapes=True)
    ```

Args:
        model: A Keras model instance
        to_file: File name of the plot image.
        show_shapes: whether to display shape information.
        show_dtype: whether to display layer dtypes.
        show_layer_names: whether to display layer names.
        rankdir: `rankdir` argument passed to PyDot,
            a string specifying the format of the plot: `"TB"`
            creates a vertical plot; `"LR"` creates a horizontal plot.
        expand_nested: whether to expand nested Functional models
            into clusters.
        dpi: Image resolution in dots per inch.
        show_layer_activations: Display layer activations (only for layers that
            have an `activation` property).
        show_trainable: whether to display if a layer is trainable.

Returns:
        A Jupyter notebook Image object if Jupyter is installed.
        This enables in-line display of the model plots in notebooks.


def get_test_data(train_samples, test_samples, input_shape, num_classes, random_seed=None):
    np.random.seed(random_seed)
    total_samples = train_samples + test_samples
    samples_per_class = total_samples // num_classes
    y = np.array([i for i in range(num_classes) for _ in range(samples_per_class)], dtype=np.int32)
    extra_samples = total_samples - len(y)
    y_extra = np.array([i % num_classes for i in range(extra_samples)], dtype=np.int64)
    y = np.concatenate([y, y_extra])
    templates = 2 * num_classes * np.random.random((num_classes,) + input_shape)
    x = np.zeros((total_samples,) + input_shape, dtype=np.float32)
    for i in range(total_samples):
        x[i] = templates[y[i]] + np.random.normal(loc=0, scale=1.0, size=input_shape)
    indices = np.arange(total_samples)
    np.random.shuffle(indices)
    x, y = (x[indices], y[indices])
    x_train, y_train, x_test, y_test = ([], [], [], [])
    for cls in range(num_classes):
        cls_indices = np.where(y == cls)[0]
        np.random.shuffle(cls_indices)
        train_count = int(train_samples / num_classes)
        x_train.extend(x[cls_indices[:train_count]])
        y_train.extend(y[cls_indices[:train_count]])
        x_test.extend(x[cls_indices[train_count:]])
        y_test.extend(y[cls_indices[train_count:]])
    x_train, y_train = (np.array(x_train), np.array(y_train))
    x_test, y_test = (np.array(x_test), np.array(y_test))
    train_indices = np.arange(len(x_train))
    test_indices = np.arange(len(x_test))
    np.random.shuffle(train_indices)
    np.random.shuffle(test_indices)
    x_train, y_train = (x_train[train_indices], y_train[train_indices])
    x_test, y_test = (x_test[test_indices], y_test[test_indices])
    return ((x_train, y_train), (x_test, y_test))

Generates balanced, stratified synthetic test data to train a model on.

Args:
        train_samples: Integer, how many training samples to generate.
        test_samples: Integer, how many test samples to generate.
        input_shape: Tuple of integers, shape of the inputs.
        num_classes: Integer, number of classes for the data and targets.
        random_seed: Integer, random seed used by Numpy to generate data.

Returns:
        A tuple of Numpy arrays: `(x_train, y_train), (x_test, y_test)`.


def _select_class_id(ids, selected_id):
    ids = sparse_tensor.convert_to_tensor_or_sparse_tensor(ids)
    if isinstance(ids, sparse_tensor.SparseTensor):
        return sparse_ops.sparse_retain(ids, math_ops.equal(ids.values, selected_id))
    ids_shape = array_ops.shape(ids, out_type=dtypes.int64)
    ids_last_dim = array_ops.size(ids_shape) - 1
    filled_selected_id_shape = math_ops.reduced_shape(ids_shape, array_ops.reshape(ids_last_dim, [1]))
    filled_selected_id = array_ops.fill(filled_selected_id_shape, math_ops.cast(selected_id, dtypes.int64))
    result = sets.set_intersection(filled_selected_id, ids)
    return sparse_tensor.SparseTensor(indices=result.indices, values=result.values, dense_shape=ids_shape)

Filter all but `selected_id` out of `ids`.

Args:
        ids: `int64` `Tensor` or `SparseTensor` of IDs.
        selected_id: Int id to select.

Returns:
        `SparseTensor` of same dimensions as `ids`. This contains only the entries
        equal to `selected_id`.


def add_model_tags(self, tags: Union[List[str], str]) -> None:
    if isinstance(tags, str):
        tags = [tags]
    if self.model_tags is None:
        self.model_tags = []
    for tag in tags:
        if tag not in self.model_tags:
            self.model_tags.append(tag)

Add custom tags into the model that gets pushed to the Hugging Face Hub. Will
    not overwrite existing tags in the model.

Args:
        tags (`Union[List[str], str]`):
            The desired tags to inject in the model

Example:
    ```python
    from transformers import AutoModel

    model = AutoModel.from_pretrained("google-bert/bert-base-cased")

    model.add_model_tags(["custom", "custom-bert"])

    # Push the model to your namespace with the name "my-custom-bert".
    model.push_to_hub("my-custom-bert")
    ```


def squareplus(x, b=4):
    if any_symbolic_tensors((x,)):
        return Squareplus(b).symbolic_call(x)
    return backend.nn.squareplus(x, b)

Squareplus activation function.

    The Squareplus activation function is defined as:

    `f(x) = (x + sqrt(x^2 + b)) / 2`

Args:
        x: Input tensor.
        b: Smoothness parameter. Defaults to 4.

Returns:
        A tensor with the same shape as `x`.

Example:
    >>> x = np.array([-1.0, 0.0, 1.0])
    >>> x_squareplus = keras.ops.squareplus(x)
    >>> print(x_squareplus)
    array([0.6180, 1.0000, 1.6180], dtype=float32)


def emit_counter(self, category: str, name: str, pid: int, timestamp: int, counter: str, value: int) -> None:
    event = self._create_event('C', category, name, pid, 0, timestamp)
    event['args'] = {counter: value}
    self._events.append(event)

Emits a record for a single counter.

Args:
        category: The event category as a string.
        name:  The event name as a string.
        pid:  Identifier of the process generating this event as an integer.
        timestamp:  The timestamp of this event as a long integer.
        counter: Name of the counter as a string.
        value:  Value of the counter as an integer.


def replace_tensors_by_numpy_ndarrays(repr_ds: RepresentativeDataset, sess: session.Session) -> RepresentativeDataset:
    new_repr_ds = []
    for sample in repr_ds:
        new_sample = {}
        for input_key, input_data in sample.items():
            if isinstance(input_data, core.Tensor):
                input_data = input_data.eval(session=sess)
            new_sample[input_key] = input_data
        new_repr_ds.append(new_sample)
    return new_repr_ds

Replaces tf.Tensors in samples by their evaluated numpy arrays.

    Note: This should be run in graph mode (default in TF1) only.

Args:
        repr_ds: Representative dataset to replace the tf.Tensors with their
        evaluated values. `repr_ds` is iterated through, so it may not be reusable
        (e.g. if it is a generator object).
        sess: Session instance used to evaluate tf.Tensors.

Returns:
        The new representative dataset where each tf.Tensor is replaced by its
        evaluated numpy ndarrays.


def _merge_tensor_signatures(self, signatures):
    sorted_update = []
    if self._num_signature_dimensions() > 1:
        signature_indices = self._signature_types()
        for _, val in sorted(signatures.items(), key=lambda item: signature_indices[item[0]]):
            sorted_update.append(val)
        updates = array_ops_stack.stack(sorted_update, axis=0, name='merge_single_op_signatures')
    elif self._num_signature_dimensions() == 1:
        (_, val), = signatures.items()
        updates = val
    else:
        raise ValueError('Cannot merge 0 signatures. Check the value passed for flag --signatures.')
    return updates

Returns a tensor that merges the given signatures.

Args:
        signatures: A dictionary of the signature updates from signature name to
        a tensor of dimension [1].

Returns:
        A tensor that concats the signature values in a predefined order.

Raises:
        ValueError: Unable to merge signatures.


def test_encode_with_non_root_fhir_path_constraint_succeeds(self, fhir_path_expression: str, expected_sql_expression: str, expected_fhir_path_sql_expression: str, expected_fields_referenced: List[str]):
    self.maxDiff = None
    constraint = self.build_constraint(fhir_path_expression=fhir_path_expression)
    self.assert_constraint_is_equal_to_expression(base_id='Hospital', element_definition_id='Hospital.patients', constraint=constraint, expected_sql_expression=expected_sql_expression, expected_fhir_path_sql_expression=expected_fhir_path_sql_expression, expected_fields_referenced=expected_fields_referenced)

Tests that a "transitive constraint" is properly encoded.

    A "transitive constraint" is a constraint defined relative to a resource
    elsewhere in the FHIR resource graph than what we're querying against.

Args:
        fhir_path_expression: The FHIRPath expression to encode.
        expected_sql_expression: The expected generated Standard SQL.
        expected_fhir_path_sql_expression: The expected generated Standard SQL
        without any contextual subqueries.
        expected_fields_referenced: The expected fields_referenced_by_expression
        attribute on the resulting constraint.


def uses_keras_history(tensors):
    checked_tensors = set()
    tensors_to_check = nest.flatten(tensors)
    while tensors_to_check:
        new_tensors_to_check = []
        for tensor in tensors_to_check:
            if id(tensor) in checked_tensors:
                continue
            checked_tensors.add(id(tensor))
            if getattr(tensor, '_keras_history_checked', None) is not None:
                continue
            if getattr(tensor, '_keras_history', None) is not None:
                return True
            try:
                new_tensors_to_check.extend(tensor.op.inputs)
            except AttributeError:
                pass
        tensors_to_check = new_tensors_to_check
    mark_checked(tensors)
    return False

Check if at least one Tensor originates from a `keras.Input`.

    This is `True` if at least one Tensor has its origin in a `keras.Input`.
    Any Tensor that originates from a `keras.Input` will have a dependency
    Tensor with a `_keras_history` attribute attached. Tensors that have
    already been checked to not originate from a `keras.Input`
    are marked as `_keras_history_checked`.

Args:
        tensors: An arbitrary nested structure of Tensors.

Returns:
        Bool, whether at least one Tensor originates from a `keras.Input`.