def session(self, graph=None, config=None, target=None):
    config = self._create_config(config)
    if target is None:
        target = self._default_target
    with session.Session(graph=graph, config=config, target=target) as sess:
        yield sess

Create a test session with master target set to the testing cluster.

    Creates a test session that connects to the local testing cluster.

Args:
        graph: Optional graph to use during the returned session.
        config: An optional config_pb2.ConfigProto to use to configure the
        session.
        target: the target of session to connect to.

Yields:
        A Session object that should be used as a context manager to surround
        the graph building and execution code in a test case.


def convert_variables_to_constants_v2(func, lower_control_flow=True, aggressive_inlining=False):
    converter_data = _FunctionConverterDataInEager(func=func, lower_control_flow=lower_control_flow, aggressive_inlining=aggressive_inlining)
    output_graph_def, converted_input_indices = _replace_variables_by_constants(converter_data=converter_data)
    return _construct_concrete_function(func, output_graph_def, converted_input_indices)

Replaces all the variables in a graph with constants of the same values.

    TensorFlow 2.0 function for converting all Variable ops into Const ops holding
    the same values. This makes it possible to describe the network fully with a
    single GraphDef file, and allows the removal of a lot of ops related to
    loading and saving the variables. This function runs Grappler's function
    inlining optimization in order to return a single subgraph.

    The current implementation only works for graphs that do not contain any
    control flow or embedding related ops.

Args:
        func: ConcreteFunction.
        lower_control_flow: Boolean indicating whether or not to lower control flow
        ops such as If and While. (default True)
        aggressive_inlining: Boolean indicating whether or not to do aggressive
        function inlining (might be unsafe if function has stateful ops, not
        properly connected to control outputs). (default False)

Returns:
        ConcreteFunction containing a simplified version of the original.


def convert_ids_to_tokens(self, ids: Union[int, list[int]], skip_special_tokens: bool=False) -> Union[str, list[str]]:
    if isinstance(ids, int):
        return self._tokenizer.id_to_token(ids)
    tokens = []
    ids_to_skip = set(self.all_special_ids) if skip_special_tokens else set()
    for index in ids:
        index = int(index)
        if index in ids_to_skip:
            continue
        tokens.append(self._tokenizer.id_to_token(index))
    return tokens

Converts a single index or a sequence of indices in a token or a sequence of tokens, using the vocabulary and
    added tokens.

Args:
        ids (`int` or `List[int]`):
            The token id (or token ids) to convert to tokens.
        skip_special_tokens (`bool`, *optional*, defaults to `False`):
            Whether or not to remove special tokens in the decoding.

Returns:
        `str` or `List[str]`: The decoded token(s).


def _create_variables_and_slots(self) -> Dict[Text, Dict[Text, tf_variables.Variable]]:
    variables = {}
    for table in self._table_config:
        variables[table.name] = self._create_variables(table, trainable=True)
    return variables

Create variables for TPU embeddings.

    Note that this will always ensure that the variable is created under the
    TPUStrategy.

Returns:
        A dict of dicts. The outer dict is keyed by the table names and the inner
        dicts are keyed by 'parameters' and the slot variable names.


def __init__(self, qualifier, args, func_set=None):
    Value.__init__(self, None)
    self.qualifier = qualifier
    self.qual_type = qualifier.qual_type
    self.args = args
    self.func_set = func_set

Resolved value from a QualifierValue.

    This is resolved instance of a QualifierValue; it is used directly by the
    test driver to validate and generate field values. This class accepts a
    list of resolved arguments which can be python primitives, STL local
    variables, or runnable STL functions.

Attributes:
        qualifier: qualifier.Qualifier instance.
        qual_type: The field type to be qualified (int, bool, string, message).
        args: List of concrete arguments to be run in
        self.qualifier.Validate()/Generate().
        func_set: Output FuncSet for this qualifier. Qualifiers can be used to
        assign values to local variables in a Transition or to fields in a
        Role via func_set.


def configure_tests(tests, test_run_id):
    print('UPDATE CONFIG')
    os.makedirs(HARNESS_DIRECTORY, exist_ok=True)
    for filename, script in tests:
        script_fields = json_get_fields(script)
        script_name = filename.split('.')[0]
        harness_fields = {}
        harness_path = HARNESS_DIRECTORY + script_name + '.json'
        if os.path.exists(harness_path):
            with open(harness_path, 'r') as f:
                harness_fields = json.load(f)
        new_fields = {}
        for field in script_fields:
            if field['name'] == 'test_run_id':
                new_fields['test_run_id'] = test_run_id
            else:
                new_fields[field['name']] = harness_fields.get(field['name'], field.get('default'))
                new_fields['%s_description' % field['name']] = '(%s) %s' % (field.get('kind', 'string'), field.get('description', 'No description.'))
                if field['name'] not in harness_fields:
                    print('NEW FIELD ADDED', script_name, field['name'])
        if new_fields:
            with open(harness_path, 'w') as f:
                json.dump(new_fields, f, indent=2)
        elif os.path.exists(harness_path):
            os.remove(harness_path)
    print('')
    print('------')
    print('------------')
    print('------------------------')
    print('Some tests require custom values. Update the necessary fields for the tests you wish to run.')
    print('EDIT: ' + HARNESS_DIRECTORY)
    print('------------------------')
    print('Some tests require external assets.  Join the following group to gain access.')
    print('VISIT: https://groups.google.com/forum/#!forum/starthinker-assets')
    print('------------------------')
    print('------------')
    print('------')
    print('')
    sleep(3)

Initialize the starthinker_assets/tests.json variable harness.

    Read all existing tests from tests/*.json and create a harness file in
    starthinker_assets/tests/*.json so developer can configure tests.

Args:
        test: List of (filename, json) pairs containing all the tests.

Returns:
        None


def make_pool3d_tests(pool_op):

    def f(options, expected_tf_failures=0):
        """Actual function that generates examples.

    Args:
      options: An Options instance.
      expected_tf_failures: number of expected tensorflow failures.
    """
        test_parameters = [{'ksize': [[1, 1, 1, 1, 1], [1, 2, 2, 2, 1], [1, 2, 3, 4, 1]], 'strides': [[1, 1, 1, 1, 1], [1, 2, 1, 2, 1], [1, 2, 2, 4, 1]], 'input_shape': [[1, 1, 1, 1, 1], [1, 16, 15, 14, 1], [3, 16, 15, 14, 3]], 'padding': ['SAME', 'VALID'], 'data_format': ['NDHWC']}]

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

        def build_inputs(parameters, sess, inputs, outputs):
            input_values = create_tensor_data(tf.float32, parameters['input_shape'])
            return ([input_values], sess.run(outputs, feed_dict=dict(zip(inputs, [input_values]))))
        extra_convert_options = ExtraConvertOptions()
        extra_convert_options.allow_custom_ops = True
        make_zip_of_tests(options, test_parameters, build_graph, build_inputs, extra_convert_options, expected_tf_failures=expected_tf_failures)
    return f

Make a set of tests to do pooling.

Args:
        pool_op: TensorFlow pooling operation to test  i.e. `tf.nn.max_pool3d`.

Returns:
        A function representing the true generator (after curried pool_op).


def builder(structdef_url: str, fhir_context: context.FhirPathContext) -> expressions.Builder:
    structdef = fhir_context.get_structure_definition(structdef_url)
    struct_type = _fhir_path_data_types.StructureDataType.from_proto(structdef)
    return expressions.Builder(_evaluation.RootMessageNode(fhir_context, struct_type), _PRIMITIVE_HANDLER)

Returns a FHIRPath expression builder.

    This gives the caller tab suggestions and early error detection when
    building FHIRPath expressions. See the documentation on the returned
    expressions.Builder for more details.

Args:
        structdef_url: the URL of the FHIR StructureDefinition to use.
        fhir_context: a DefinitionLoader used to load FHIR structure definitions and
        dependencies.
    Returns: a builder object to creae FHIRPath expressions.


def index_of(self, file_path, function_name, function_start_line):
    function_key = (file_path, function_name, function_start_line)
    if function_key in self._function_key_to_function:
        return self._function_key_to_function[function_key].id
    else:
        function_index = len(self._function_key_to_function) + 1
        function = profile_pb2.Function()
        function.id = function_index
        function.name = self._string_table.index_of(function_name)
        function.filename = self._string_table.index_of(file_path)
        function.start_line = function_start_line
        self._function_key_to_function[function_key] = function
        return function_index

Returns index of the function, adding the function if needed.

Args:
        file_path: (string) Path to file where the function is defined.
        function_name: (string) Function name.
        function_start_line: (integer) Start line number of function definition.

Returns:
        Function index.


def bind_object(self, trackable):
    checkpoint = self.checkpoint
    checkpoint.all_python_objects.add(trackable)
    current_assignment = checkpoint.object_by_proto_id.get(self._proto_id, None)
    checkpoint.matched_proto_ids.add(self._proto_id)
    if current_assignment is None:
        checkpoint.object_by_proto_id[self._proto_id] = trackable
        return True
    else:
        if current_assignment is not trackable:
            logging.warning(f'Inconsistent references when loading the checkpoint into this object graph. For example, in the saved checkpoint object, `model.layer.weight` and `model.layer_copy.weight` reference the same variable, while in the current object these are two different variables. The referenced variables are:({current_assignment} and {trackable}).')
        return False

Set a checkpoint<->object correspondence.

Args:
        trackable: The object to record a correspondence for.

Returns:
        True if this is a new assignment, False if this object has already been
        mapped to a checkpointed `Object` proto.

Raises:
        AssertionError: If another object is already bound to the `Object` proto.


def compute_pooling_output_shape(input_shape, pool_size, strides, padding='valid', data_format='channels_last'):
    strides = pool_size if strides is None else strides
    input_shape_origin = list(input_shape)
    input_shape = np.array(input_shape)
    if data_format == 'channels_last':
        spatial_shape = input_shape[1:-1]
    else:
        spatial_shape = input_shape[2:]
    none_dims = []
    for i in range(len(spatial_shape)):
        if spatial_shape[i] is None:
            spatial_shape[i] = -1
            none_dims.append(i)
    pool_size = np.array(pool_size)
    if padding == 'valid':
        output_spatial_shape = np.floor((spatial_shape - pool_size) / strides) + 1
        for i in range(len(output_spatial_shape)):
            if i not in none_dims and output_spatial_shape[i] < 0:
                raise ValueError(f'Computed output size would be negative. Received: `inputs.shape={input_shape}` and `pool_size={pool_size}`.')
    elif padding == 'same':
        output_spatial_shape = np.floor((spatial_shape - 1) / strides) + 1
    else:
        raise ValueError(f"Argument `padding` must be either 'valid' or 'same'. Received: padding={padding}")
    output_spatial_shape = [int(i) for i in output_spatial_shape]
    for i in none_dims:
        output_spatial_shape[i] = None
    output_spatial_shape = tuple(output_spatial_shape)
    if data_format == 'channels_last':
        output_shape = (input_shape_origin[0],) + output_spatial_shape + (input_shape_origin[-1],)
    else:
        output_shape = (input_shape_origin[0], input_shape_origin[1]) + output_spatial_shape
    return output_shape

Computes the output shape of pooling operations.

Args:
        input_shape: Input shape. Must be a tuple of integers.
        pool_size: Size of the pooling operation. Must be a tuple of integers.
        strides: Stride of the pooling operation. Must be a tuple of integers.
            Defaults to `pool_size`.
        padding: Padding method. Available methods are `"valid"` or `"same"`.
            Defaults to `"valid"`.
        data_format: String, either `"channels_last"` or `"channels_first"`.
            The ordering of the dimensions in the inputs. `"channels_last"`
            corresponds to inputs with shape `(batch, height, width, channels)`
            while `"channels_first"` corresponds to inputs with shape
            `(batch, channels, height, weight)`. Defaults to `"channels_last"`.

Returns:
        Tuple of ints: The output shape of the pooling operation.

Example:
    # Basic usage with square pooling on a single image
    >>> compute_pooling_output_shape((1, 4, 4, 1), (2, 2))
    (1, 2, 2, 1)

    # Strided pooling on a single image with strides different from pool_size
    >>> compute_pooling_output_shape((1, 4, 4, 1), (2, 2), strides=(1, 1))
    (1, 3, 3, 1)

    # Pooling on a batch of images
    >>> compute_pooling_output_shape((32, 4, 4, 3), (2, 2))
    (32, 2, 2, 3)


def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader:
    return self._get_dataloader(dataset=test_dataset, description='test', batch_size=self.args.eval_batch_size, sampler_fn=self._get_eval_sampler)

Returns the test [`~torch.utils.data.DataLoader`].

    Subclass and override this method if you want to inject some custom behavior.

Args:
        test_dataset (`torch.utils.data.Dataset`, *optional*):
            The test dataset to use. If it is a [`~datasets.Dataset`], columns not accepted by the
            `model.forward()` method are automatically removed. It must implement `__len__`.


def apply_sync(fn: StreamFn, content: Iterable[_T]) -> list[_T]:

    async def run_with_context():
        async with context.context():
            as_async = streams.stream_content(content)
            return await streams.gather_stream(fn(as_async))
    return asyncio.run(run_with_context())

Applies a part function synchronously.

Args:
        fn: the part function to apply to the content.
        content: a collection of inputs/parts on which to apply the function.

Returns:
        the content, with the function `fn` applied to each input/part.


def scalar_spec(value_spec: pg.typing.ValueSpec) -> pg.typing.ValueSpec:
    return pg.typing.Union([value_spec, pg.typing.Callable([pg.typing.Int()], returns=value_spec)])

Returns the value spec for a schedule scalar.

Args:
        value_spec: a value spec for the schedule-based scalar type.

Returns:
        A value spec for either the value itself or a callable that produces such
        value based on a step (integer).


def resize_position_embeddings(self, new_num_position_embeddings: int):
    self.config.max_position_embeddings = new_num_position_embeddings
    self.encoder.resize_position_embeddings(new_num_position_embeddings)
    self.decoder.resize_position_embeddings(new_num_position_embeddings)

Resizes position embeddings matrix of the model if `new_num_position_embeddings !=
    config.max_position_embeddings`.

Args:
        new_num_position_embeddings (`int`):
            The number of new position embeddings. If position embeddings are learned, increasing the size will add
            newly initialized vectors at the end, whereas reducing the size will remove vectors from the end. If
            position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the size will
            add correct vectors at the end following the position encoding algorithm, whereas reducing the size
            will remove vectors from the end.


def __init__(self, name='squared_hinge', dtype=None):
    super(SquaredHinge, self).__init__(squared_hinge, name, dtype=dtype)

Computes the squared hinge metric between `y_true` and `y_pred`.

    `y_true` values are expected to be -1 or 1. If binary (0 or 1) labels are
    provided we will convert them to -1 or 1.

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

    Standalone usage:

    >>> m = tf.keras.metrics.SquaredHinge()
    >>> m.update_state([[0, 1], [0, 0]], [[0.6, 0.4], [0.4, 0.6]])
    >>> m.result().numpy()
    1.86

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

    Usage with `compile()` API:

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


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

Create a mask from the two sequences passed to be used in a sequence-pair classification task. XLM-RoBERTa does
    not make use of token type ids, therefore a list of zeros is returned.

Args:
        token_ids_0 (`List[int]`):
            List of IDs.
        token_ids_1 (`List[int]`, *optional*):
            Optional second list of IDs for sequence pairs.

Returns:
        `List[int]`: List of zeros.


def merge_bindings(program: cfg.Program, node: cfg.CFGNode, bindings: Sequence[cfg.Binding]) -> cfg.Variable:
    v = program.NewVariable()
    for b in bindings:
        v.PasteBinding(b, node)
    return v

Create a combined Variable for a list of bindings.

Args:
        program: A cfg.Program instance.
        node: The current CFG node.
        bindings: A list of cfg.Bindings.

Returns:
        A cfg.Variable.


def cluster_spec(self):
    if self._tpu != 'local':
        network_endpoints = self._cloud_tpu_client.network_endpoints()
        worker_list = ['%s:%s' % (endpoint['ipAddress'], endpoint['port']) for endpoint in network_endpoints]
        cluster_spec = {self.task_type: worker_list}
        if self._coordinator_address:
            cluster_spec[self._coordinator_name] = [self._coordinator_address]
        return server_lib.ClusterSpec(cluster_spec)
    else:
        return server_lib.ClusterSpec({})

Returns a ClusterSpec object based on the latest TPU information.

    We retrieve the information from the GCE APIs every time this method is
    called.

Returns:
        A ClusterSpec containing host information returned from Cloud TPUs,
        or None.

Raises:
        RuntimeError: If the provided TPU is not healthy.


def _scoped_subscribe(tensor, side_effects, control_cache):
    with ops.device(tensor.device):
        with _preserve_control_flow_context(tensor):
            return _subscribe(tensor, side_effects, control_cache)

Helper method that subscribes a single tensor to a list of side_effects.

    This is a thin wrapper around `_subscribe` and ensures that the side effect
    ops are added within the same device and control flow context of the
    subscribed tensor.

Args:
        tensor: The `tf.Tensor` to be subscribed.
        side_effects: List of side_effect functions, see subscribe for details.
        control_cache: `_ControlOutputCache` helper to get control_outputs faster.

Returns:
        The modified replacement to the passed in tensor which triggers the side
        effects or the given tensor, if it was already been subscribed.


def call(self, pixel_values: tf.Tensor, input_ids: 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, TFBlipForConditionalGenerationModelOutput]:
    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]
    outputs = self.text_decoder(input_ids=input_ids, attention_mask=attention_mask, encoder_hidden_states=image_embeds, labels=labels, return_dict=False, training=training)
    if not return_dict:
        outputs = (outputs[0], outputs[1], image_embeds, vision_outputs[0]) + vision_outputs[2:]
        return tuple((output for output in outputs if output is not None))
    if labels is not None:
        loss = outputs[0]
        logits = outputs[1]
    else:
        loss = None
        logits = outputs[0]
    if loss is not None and loss.shape.rank == 0:
        loss = tf.reshape(loss, (1,))
    return TFBlipForConditionalGenerationModelOutput(loss=loss, logits=logits, 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, TFBlipForConditionalGeneration

    >>> processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
    >>> model = TFBlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

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

    >>> inputs = processor(images=image, text=text, return_tensors="tf")

    >>> outputs = model(**inputs)
    ```


def _get_clang_major_version(path_to_clang: str) -> int:
    logging.info('Running echo __clang_major__ | %s -E -P -', path_to_clang)
    clang_version_proc = subprocess.run([path_to_clang, '-E', '-P', '-'], input='__clang_major__', check=True, capture_output=True, text=True)
    major_version = int(clang_version_proc.stdout)
    logging.info('%s reports major version %s.', path_to_clang, major_version)
    return major_version

Gets the major version of the clang at `path_to_clang`.

Args:
        path_to_clang: Path to a clang executable

Returns:
        The major version.


def __init__(self, engine: trt.ICudaEngine):
    from cuda import cuda
    import tensorrt as trt
    self.engine = engine
    self.context = engine.create_execution_context()
    self.context_lock = threading.RLock()
    self.inputs = []
    self.outputs = []
    self.gpu_allocations = []
    self.cpu_allocations = []
    try:
        _ = np.bool
    except AttributeError:
        np.bool = np.bool_
    for i in range(self.engine.num_bindings):
        name = self.engine.get_binding_name(i)
        dtype = self.engine.get_binding_dtype(i)
        shape = self.engine.get_binding_shape(i)
        size = trt.volume(shape) * dtype.itemsize
        allocation = _assign_or_fail(cuda.cuMemAlloc(size))
        binding = {'index': i, 'name': name, 'dtype': np.dtype(trt.nptype(dtype)), 'shape': list(shape), 'allocation': allocation, 'size': size}
        self.gpu_allocations.append(allocation)
        if self.engine.binding_is_input(i):
            self.inputs.append(binding)
        else:
            self.outputs.append(binding)
    assert self.context
    assert len(self.inputs) > 0
    assert len(self.outputs) > 0
    assert len(self.gpu_allocations) > 0
    for output in self.outputs:
        self.cpu_allocations.append(np.zeros(output['shape'], output['dtype']))
    self.stream = _assign_or_fail(cuda.cuStreamCreate(0))

Implementation of the TensorRTEngine class which handles
    allocations associated with TensorRT engine.

    Example Usage::

        TensorRTEngine(engine)

Args:
        engine: trt.ICudaEngine object that contains TensorRT engine


def __init__(self, detector: AnomalyDetector):
    self._detector = detector

Applies the _ScoreAndLearnDoFn to a PCollection of data.

    This PTransform scores and learns from data points using an anomaly
    detection model.

Args:
        detector: The anomaly detection model to use.


def call(self, input_ids: tf.Tensor | None=None, pixel_values: tf.Tensor | None=None, attention_mask: tf.Tensor | None=None, position_ids: tf.Tensor | None=None, return_loss: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, training: Optional[bool]=None) -> Union[Tuple, TFBlipOutput]:
    outputs = self.blip(input_ids=input_ids, pixel_values=pixel_values, attention_mask=attention_mask, position_ids=position_ids, return_loss=return_loss, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, training=training)
    return outputs

Returns:

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

    >>> model = TFBlipModel.from_pretrained("Salesforce/blip-image-captioning-base")
    >>> processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")

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

    >>> inputs = processor(
    ...     text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="tf", padding=True
    ... )

    >>> outputs = model(**inputs)
    >>> logits_per_image = outputs.logits_per_image  # this is the image-text similarity score
    >>> probs = tf.nn.softmax(logits_per_image, axis=1)  # we can take the softmax to get the label probabilities
    ```


def _tf_glue_convert_examples_to_features(examples: tf.data.Dataset, tokenizer: PreTrainedTokenizer, task=str, max_length: Optional[int]=None) -> tf.data.Dataset:
    processor = glue_processors[task]()
    examples = [processor.tfds_map(processor.get_example_from_tensor_dict(example)) for example in examples]
    features = glue_convert_examples_to_features(examples, tokenizer, max_length=max_length, task=task)
    label_type = tf.float32 if task == 'sts-b' else tf.int64

    def gen():
        for ex in features:
            d = {k: v for k, v in asdict(ex).items() if v is not None}
            label = d.pop('label')
            yield (d, label)
    input_names = tokenizer.model_input_names
    return tf.data.Dataset.from_generator(gen, (dict.fromkeys(input_names, tf.int32), label_type), ({k: tf.TensorShape([None]) for k in input_names}, tf.TensorShape([])))

Returns:
    A `tf.data.Dataset` containing the task-specific features.


def create(self, url, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO) -> BinaryIO:
    _, path = self._parse_url(url)
    return self._create(path, mime_type, compression_type)


def kaiser(x, beta):
    if any_symbolic_tensors((x,)):
        return Kaiser(beta).symbolic_call(x)
    return backend.numpy.kaiser(x, beta)

Kaiser window function.

    The Kaiser window is defined as:
    `w[n] = I0(beta * sqrt(1 - (2n / (N - 1) - 1)^2)) / I0(beta)`
    where I0 is the modified zeroth-order Bessel function of the first kind.

Args:
        x: Scalar or 1D Tensor. The window length.
        beta: Float. Shape parameter for the Kaiser window.

Returns:
        A 1D tensor containing the Kaiser window values.

Example:
    >>> x = keras.ops.convert_to_tensor(5)
    >>> keras.ops.kaiser(x, beta=14.0)
    array([7.7268669e-06, 1.6493219e-01, 1.0000000e+00, 1.6493219e-01,
        7.7268669e-06], dtype=float32)


def assert_same_structure(nest1, nest2, check_types=True):
    nest_util.assert_same_structure(nest_util.Modality.DATA, nest1, nest2, check_types)

Asserts that two structures are nested in the same way.

Args:
        nest1: an arbitrarily nested structure.
        nest2: an arbitrarily nested structure.
        check_types: if `True` (default) types of sequences should be same as
        well. For dictionary, "type" of dictionary is considered to include its
        keys. In other words, two dictionaries with different keys are considered
        to have a different "type". If set to `False`, two iterables are
        considered same as long as they yield the elements that have same
        structures.

Raises:
        ValueError: If the two structures do not have the same number of elements or
        if the two structures are not nested in the same way.
        TypeError: If the two structures differ in the type of sequence in any of
        their substructures. Only possible if `check_types` is `True`.


def _rename_if_any_arg_found_transformer(parent, node, full_name, name, logs, arg_names=None, arg_ok_predicate=None, remove_if_ok=False, message=None):
    for arg_name in arg_names:
        rename_node = _rename_if_arg_found_transformer(parent, node, full_name, name, logs, arg_name, arg_ok_predicate, remove_if_ok, message)
        node = rename_node if rename_node else node
    return node

Replaces the given call with tf.compat.v1 if any of the arg_names is found.

Args:
        parent: Parent of node.
        node: ast.Call node to modify.
        full_name: full name of function to modify.
        name: name of function to modify.
        logs: list of logs to append to.
        arg_names: list of names of the argument to look for.
        arg_ok_predicate: predicate callable with the ast of the argument value,
        returns whether the argument value is allowed.
        remove_if_ok: remove the argument if present and ok as determined by
        arg_ok_predicate.
        message: message to print if a non-ok arg is found (and hence, the function
        is renamed to its compat.v1 version).

Returns:
        node, if it was modified, else None.


Called at the end of prediction.

    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 __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(AveragePooling3D, self).__init__(pool_size=pool_size, strides=strides, padding=padding, data_format=data_format, name=name, **kwargs)

Average 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 all_reduce(self, input_tensor: core.TensorLike, control_input: Optional[Union[core.TensorLike, ops.Operation]]=None, options: Optional[collective_util.Options]=None) -> core.Tensor:
    instance_key = self._next_instance_key()
    options = self._options.merge(options)
    ordering_token = self._get_ordering_token()
    with ops.device(self._device), self._control_input(control_input):
        return collective_ops.all_reduce_v2(input_tensor, self._group_size, self._group_key, instance_key, communication_hint=options.implementation.value, timeout=options.timeout_seconds, ordering_token=ordering_token)

All-reduce a dense tensor.

Args:
        input_tensor: a dense tensor. It must have the same shape on all replicas.
        control_input: if not None, add control edges between control_input and
        the all-reduce.
        options: an optional tf.distribute.experimental.CommunicationOptions. If
        provided, it overrides the default options.

Returns:
        The reduced tensor.


def _name_list(tensor_list):
    return [compat.as_bytes(t.name) for t in tensor_list]

Utility function for transitioning to the new session API.

Args:
        tensor_list: a list of `Tensor`s.

Returns:
        A list of each `Tensor`s name (as byte arrays).


def load_model(file_path: str) -> NormalizedModel:
    with open(file_path) as f:
        model = json.load(f)
    model_flat = OrderedDict()
    for category in model:
        for item in model[category]:
            model_flat['%s:%s' % (category, item)] = model[category][item]
    weights = jnp.array(list(model_flat.values()))
    weights = weights / weights.std()
    weights = weights - weights.mean()
    keys = list(model_flat.keys())
    return NormalizedModel(keys, weights)

Loads a model as a pair of a features list and a normalized weight vector.

Args:
        file_path: A file path for the model JSON file.

Returns:
        A normalized model, which is a pair of a list of feature identifiers and a
        normalized weight vector.


def load_function_def_library(library, saved_object_graph=None, load_shared_name_suffix=None, wrapper_function=None):
    library_function_names = set((fdef.signature.name for fdef in library.function))
    functions = {}
    renamed_functions = {}
    if ops.executing_eagerly_outside_functions():
        graph = ops.Graph()
    else:
        graph = ops.get_default_graph()
    if load_shared_name_suffix is None:
        load_shared_name_suffix = '_load_{}'.format(ops.uid())
    library_gradient_names = {}
    new_gradient_op_types = {}
    gradients_to_register = {}
    for gdef in library.registered_gradients:
        if gdef.registered_op_type:
            new_op_type = custom_gradient.generate_name()
            old_op_type = compat.as_bytes(gdef.registered_op_type)
            library_gradient_names[old_op_type] = gdef.gradient_func
            new_gradient_op_types[old_op_type] = new_op_type
            gradients_to_register[gdef.gradient_func] = new_op_type
    function_deps = {}
    for fdef in library.function:
        function_deps[fdef.signature.name] = _list_function_deps(fdef, library_function_names, library_gradient_names)
    loaded_gradients = {}
    for fdef in _sort_function_defs(library, function_deps):
        orig_name = _fix_fdef_in_place(fdef, functions, load_shared_name_suffix, new_gradient_op_types)
        structured_input_signature = None
        structured_outputs = None
        if saved_object_graph is not None and orig_name in saved_object_graph.concrete_functions:
            proto = saved_object_graph.concrete_functions[orig_name]
            structured_input_signature = nested_structure_coder.decode_proto(proto.canonicalized_input_signature)
            structured_outputs = nested_structure_coder.decode_proto(proto.output_signature)
        with graph.as_default():
            func_graph = function_def_lib.function_def_to_graph(fdef, structured_input_signature=structured_input_signature, structured_outputs=structured_outputs)
        _restore_gradient_functions(func_graph, renamed_functions, loaded_gradients)
        for dep in function_deps[orig_name]:
            functions[dep].add_to_graph(func_graph)
        if '_input_shapes' in fdef.attr:
            del fdef.attr['_input_shapes']
        function_type = function_type_lib.from_structured_signature(func_graph.structured_input_signature, func_graph.structured_outputs, func_graph.function_captures.capture_types)
        func = function_lib.ConcreteFunction.from_func_graph(func_graph, function_type, attrs=fdef.attr)
        if wrapper_function:
            func = wrapper_function(func)
        func.add_to_graph(graph)
        functions[orig_name] = func
        renamed_functions[func.name] = func
        if any((op.type == 'TRTEngineOp' for op in func_graph.get_operations())):
            func.add_to_graph(ops.get_default_graph())
        if orig_name in gradients_to_register:
            gradient_op_type = gradients_to_register[orig_name]
            loaded_gradients[compat.as_bytes(gradient_op_type)] = func
            ops.RegisterGradient(gradient_op_type)(_gen_gradient_func(func))
    return functions

Load a set of functions as concrete functions without captured inputs.

    Functions names are manipulated during load such that they do not overlap
    with previously created ones.

    Gradients are re-registered under new names. Ops that reference the gradients
    are updated to reflect the new registered names.

Args:
        library: FunctionDefLibrary proto message.
        saved_object_graph: SavedObjectGraph proto message. If not passed in,
        concrete function structured signatures and outputs will not be set.
        load_shared_name_suffix: If specified, used to uniquify shared names.
        Otherwise, a unique name is generated.
        wrapper_function: An object that will be wrapped on newly created functions.

Returns:
        Map of original function names in the library to instances of
        `ConcreteFunction` without captured inputs.

Raises:
        ValueError: if functions dependencies have a cycle.


def __hash__(self) -> int:
    return self._text.__hash__()

Hash function.

    NOTE(daiyip): ConstStrKey shares the same hash with its text, which
    makes it easy to lookup a dict of string by an ConstStrKey object, and
    vice versa.

Returns:
        Hash code.


def run_without_tensor_float_32(description: str) -> Callable[[Callable[..., Any]], Callable[..., None]]:

    def decorator(f: Callable[..., Any]) -> Callable[..., None]:

        @functools.wraps(f)
        def decorated(*args, **kwargs):
            allowed = config.tensor_float_32_execution_enabled()
            try:
                config.enable_tensor_float_32_execution(False)
                f(*args, **kwargs)
            finally:
                config.enable_tensor_float_32_execution(allowed)
        return tf_decorator.make_decorator(f, decorated)
    return decorator

Execute test with TensorFloat-32 disabled.

    While almost every real-world deep learning model runs fine with
    TensorFloat-32, many tests use assertAllClose or similar methods.
    TensorFloat-32 matmuls typically will cause such methods to fail with the
    default tolerances.

Args:
        description: A description used for documentation purposes, describing why
        the test requires TensorFloat-32 to be disabled.

Returns:
        Decorator which runs a test with TensorFloat-32 disabled.


def compute_global_attention_representations(self, global_q, global_k, global_v, local_k, local_v, mask, dim: DimensionInfo):
    global_and_local_k = torch.cat([global_k, local_k], dim=2)
    global_and_local_v = torch.cat([global_v, local_v], dim=2)
    extended_mask = nn.functional.pad(mask, pad=(dim.global_len, 0), value=0)
    attn_weights = torch.einsum('BHGF,BHXF->BHGX', global_q, global_and_local_k)
    attn_weights = attn_weights + extended_mask[:, None, None, :]
    attn_probs = nn.functional.softmax(attn_weights, dim=-1)
    attn_probs = nn.functional.dropout(attn_probs, p=self.dropout, training=self.training)
    attn_output = torch.einsum('BHGX,BHXF->BHGF', attn_probs, global_and_local_v)
    return (attn_output, attn_probs)

Compute attention representations for global tokens.

    Global tokens will attend to both global tokens as well as all input sequence tokens. Because the input
    sequence tokens are arranged in blocks for local attention, we unblock them and compute attention.

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

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


def post_process_image_text_to_text(self, generated_outputs, skip_special_tokens=True, **kwargs):
    beginning_of_answer = self.tokenizer.convert_tokens_to_ids(BEGINNING_OF_ANSWER_STRING)
    unpadded_output_sequences = [seq[(seq == beginning_of_answer).nonzero(as_tuple=True)[0] + 1:] for seq in generated_outputs]
    max_len = max((len(seq) for seq in unpadded_output_sequences))
    padded_output_sequences = torch.full((len(unpadded_output_sequences), max_len), self.pad_token_id)
    for i, seq in enumerate(unpadded_output_sequences):
        padded_output_sequences[i, :len(seq)] = torch.tensor(seq)
    return self.batch_decode(padded_output_sequences, skip_special_tokens=skip_special_tokens, **kwargs)

Post-processes the output of `FuyuForConditionalGeneration` to only return the text output.

Args:
        generated_outputs (`torch.Tensor` or `np.ndarray`):
            The output of the model. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
            containing the token ids of the generated sequences.
        skip_special_tokens (`bool`, *optional*, defaults to `True`):
            Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method.
        **kwargs:
            Additional arguments to be passed to the tokenizer's `batch_decode method`.

Returns:
        `List[str]`: The decoded text output.


def get_func_graphs(op):

    def _get_func_graph_for_branch(name_attr_list, cached_attr_name=None):
        """Generates and returns a FuncGraph for the given branch."""
        func_graph = None
        if cached_attr_name is not None:
            func_graph = getattr(op, cached_attr_name, None)
        inputs = op.inputs[1:]
        if func_graph is None:
            input_shapes = [t.shape for t in inputs]
            func_graph = util.get_func_graph(op, input_shapes, name_attr_list.name)
        for external_t, internal_t in zip(inputs, func_graph.inputs):
            handle_data_util.copy_handle_data(external_t, internal_t)
        func_graph.function_captures.reset_captures(inputs, func_graph.inputs)
        func_graph._forward_cond = op
        return func_graph
    if op.type in ['If', 'StatelessIf']:
        return (_get_func_graph_for_branch(op.get_attr('then_branch'), '_true_graph'), _get_func_graph_for_branch(op.get_attr('else_branch'), '_false_graph'))
    elif op.type in ['Case', 'StatelessCase']:
        return [_get_func_graph_for_branch(branch_fn, '_branch_graph_{}'.format(i)) for i, branch_fn in enumerate(op.get_attr('branches'))]
    else:
        raise ValueError('Unsupported op type: {}'.format(op.type))

Returns `FuncGraph`s for the input op branches.

Args:
        op: The If or Case Operation.

Returns:
        A tuple of the `FuncGraph`s of the then_branch and else_branch (all branches
        for Case).


def _gen_data(m, k, n, nnz, row_occupied_rate, data_type, seed):
    random.seed(seed)
    np.random.seed(seed)
    occupied_rows = random.sample(range(m), int(m * row_occupied_rate))
    sparse_input_dense_shape = [m, k]
    dense_input_shape = (k, n)
    indices = []
    for _ in range(nnz):
        row = random.choice(occupied_rows)
        col = random.randint(0, k - 1)
        indices.append([row, col])

    def maybe_complex(x):
        if x.dtype.kind == 'c':
            return (x + 1j * x) / 2
        return x
    sparse_values = maybe_complex(np.random.normal(size=len(indices)).astype(data_type))
    dense_values = maybe_complex(np.random.normal(size=dense_input_shape).astype(data_type))
    sparse_input = sparse_tensor.SparseTensor(indices, sparse_values, sparse_input_dense_shape)
    dense_input = constant_op.constant(dense_values)
    return (sparse_input, dense_input)

Generate valid input data for tf.sparse.sparse_dense_matmul

    returns sparse matrix a (type SparseTensor), dense matrix b (type Tensor)

    Parameters:
        m: row count of dense version of matrix a / row count of output matrix
        k: col count of dense version of matrix a / row count of matrix b
        n: col could of matrix b / col count of output matrix
        nnz: number of non-zero elements in matrix a
        row_occupied_rate: prob that row in a has one or more non-zero element


def enumerate_dataset(start=0):

    def _apply_fn(dataset):
        return dataset.enumerate(start)
    return _apply_fn

A transformation that enumerates the elements of a dataset.

    It is similar to python's `enumerate`.
    For example:

    ```python
    # NOTE: The following examples use `{ ... }` to represent the
    # contents of a dataset.
    a = { 1, 2, 3 }
    b = { (7, 8), (9, 10) }

    # The nested structure of the `datasets` argument determines the
    # structure of elements in the resulting dataset.
    a.apply(tf.data.experimental.enumerate_dataset(start=5))
    => { (5, 1), (6, 2), (7, 3) }
    b.apply(tf.data.experimental.enumerate_dataset())
    => { (0, (7, 8)), (1, (9, 10)) }
    ```

Args:
        start: A `tf.int64` scalar `tf.Tensor`, representing the start value for
        enumeration.

Returns:
        A `Dataset` transformation function, which can be passed to
        `tf.data.Dataset.apply`.


def ragged_shape(input: ragged_tensor.Ragged, name: Optional[str]=None, out_type=dtypes.int32) -> dynamic_ragged_shape.DynamicRaggedShape:
    with ops.name_scope(name, 'RaggedShape', [input]):
        return dynamic_ragged_shape.DynamicRaggedShape.from_tensor(input, out_type)

Returns the shape of a RaggedTensor.

Args:
        input: A `RaggedTensor`
        name: A name for the operation (optional).
        out_type: dtype used to encode the shape.

Returns:
        A `tf.experimental.DynamicRaggedShape`


def prune_linear_layer(layer: nn.Linear, index: torch.LongTensor, dim: int=0) -> nn.Linear:
    index = index.to(layer.weight.device)
    W = layer.weight.index_select(dim, index).detach().clone()
    if layer.bias is not None:
        if dim == 1:
            b = layer.bias.detach().clone()
        else:
            b = layer.bias[index].detach().clone()
    new_size = list(layer.weight.size())
    new_size[dim] = len(index)
    new_layer = nn.Linear(new_size[1], new_size[0], bias=layer.bias is not None).to(layer.weight.device)
    new_layer.weight.requires_grad = False
    new_layer.weight.copy_(W.contiguous())
    new_layer.weight.requires_grad = True
    if layer.bias is not None:
        new_layer.bias.requires_grad = False
        new_layer.bias.copy_(b.contiguous())
        new_layer.bias.requires_grad = True
    return new_layer

Prune a linear layer to keep only entries in index.

    Used to remove heads.

Args:
        layer (`torch.nn.Linear`): The layer to prune.
        index (`torch.LongTensor`): The indices to keep in the layer.
        dim (`int`, *optional*, defaults to 0): The dimension on which to keep the indices.

Returns:
        `torch.nn.Linear`: The pruned layer as a new layer with `requires_grad=True`.


def register_for_auto_class(cls, auto_class='AutoTokenizer'):
    if not isinstance(auto_class, str):
        auto_class = auto_class.__name__
    import transformers.models.auto as auto_module
    if not hasattr(auto_module, auto_class):
        raise ValueError(f'{auto_class} is not a valid auto class.')
    cls._auto_class = auto_class

Register this class with a given auto class. This should only be used for custom tokenizers as the ones in the
    library are already mapped with `AutoTokenizer`.

Args:
        auto_class (`str` or `type`, *optional*, defaults to `"AutoTokenizer"`):
            The auto class to register this new tokenizer with.


def _FullUpdateFromFile(self, cache, source_map, force_write=False):
    return_val = 0
    for entry in source_map:
        if not entry.Verify():
            raise error.InvalidMap('Map is not valid. Aborting')
    if len(source_map) == 0 and (not force_write):
        raise error.EmptyMap('Source map empty during full update, aborting. Use --force-write to override.')
    return_val += cache.WriteMap(map_data=source_map, force_write=force_write)
    if return_val == 0:
        mtime = os.stat(cache.GetCacheFilename()).st_mtime
        self.log.debug('Cache filename %s has mtime %d', cache.GetCacheFilename(), mtime)
        self.WriteModifyTimestamp(mtime)
        self.WriteUpdateTimestamp()
    return return_val

Write a new map into the provided cache (overwrites).

Args:
        cache: A nss_cache.caches.Cache object.
        source_map: The map whose contents we're replacing the cache with, that is
        used for verification.
        force_write: A boolean flag forcing empty map updates when False,
        defaults to False.

Returns:
        0 if succesful, non-zero indicating number of failures otherwise.

Raises:
        EmptyMap: Update is an empty map, not raised if force_write=True.
        InvalidMap:


def __init__(self, dataset, class_weight=None, distribution=None):
    from keras.src.utils.module_utils import tensorflow as tf
    if not isinstance(dataset, (tf.data.Dataset, tf.distribute.DistributedDataset)):
        raise ValueError(f'Expected argument `dataset` to be a tf.data.Dataset. Received: {dataset}')
    if class_weight is not None:
        dataset = dataset.map(make_class_weight_map_fn(class_weight)).prefetch(tf.data.AUTOTUNE)
    if distribution is not None:
        dataset = distribution.distribute_dataset(dataset)
    self._dataset = dataset

Initialize the TFDatasetAdapter.

Args:
        dataset: The input `tf.data.Dataset` instance.
        class_weight: A map where the keys are integer class ids and values
            are the class weights, e.g. `{0: 0.2, 1: 0.6, 2: 0.3}`.
        distribution: A `keras.distribution.Distribution` instance. Used to
            shard the input dataset into per worker/process dataset
            instance.


def no_automatic_dependency_tracking_scope(obj):
    previous_value = getattr(obj, '_setattr_tracking', True)
    obj._setattr_tracking = False
    try:
        yield
    finally:
        obj._setattr_tracking = previous_value

A context that disables automatic dependency tracking when assigning attrs.

    Objects that inherit from Autotrackable automatically creates dependencies
    to trackable objects through attribute assignments, and wraps data structures
    (lists or dicts) with trackable classes. This scope may be used to temporarily
    disable this behavior. This works similar to the decorator
    `no_automatic_dependency_tracking`.

    Example usage:
    ```
    model = tf.keras.Model()
    model.arr1 = []  # Creates a ListWrapper object
    with no_automatic_dependency_tracking_scope(model):
        model.arr2 = []  # Creates a regular, untracked python list
    ```

Args:
        obj: A trackable object.

Yields:
        a scope in which the object doesn't track dependencies.


def __init__(self, columns: list[str], vocab_size: Optional[int]=None, smooth: bool=True, name: Optional[str]=None):
    super().__init__(columns)
    self.vocab_size = vocab_size
    self.smooth = smooth
    self.name = name
    self.tfidf_weight = None

This function applies a tf-idf transformation on the given columns
    of incoming data.

    TFIDF outputs two artifacts for each column: the vocabulary index and
    the tfidf weight. The vocabulary index is a mapping from the original
    vocabulary to the new vocabulary. The tfidf weight is a mapping
    from the original vocabulary to the tfidf score.

    Input passed to the TFIDF is not modified and used to calculate the
    required artifacts.

Args:
        columns: List of column names to apply the transformation.
        vocab_size: (Optional) An integer that specifies the size of the
        vocabulary. Defaults to None.

        If vocab_size is None, then the size of the vocabulary is
        determined by `tft.get_num_buckets_for_transformed_feature`.
        smooth: (Optional) A boolean that specifies whether to apply
        smoothing to the tf-idf score. Defaults to True.
        name: (Optional) A string that specifies the name of the operation.


def forward(ctx, forward_fn, *args, **kwargs):
    ctx.forward_fn = forward_fn
    ctx.save_for_backward(*args)
    try:
        output, ctx.grad_fn = forward_fn(*args, **kwargs)
    except:
        output = forward_fn(*args, **kwargs)
        ctx.grad_fn = lambda *args, **kwargs: torch.full((), float('nan'))
    return output

Forward pass computation specification.

Args:
        ctx: Context object.
        forward_fn: Function to compute forward pass.
        *args: Arguments for the forward pass.
        **kwargs: Keyword arguments for the forward pass.


def bloom_gelu_back(g: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
    x = x[0]
    tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
    ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out)
    return ff * g

gradient of tanh approximation of gelu gradient of actual gelu is: 0.5 * (1. + torch.erf(x * 0.70710678)) +
    0.3989423 * x * torch.exp(-0.5 * x * x)

Args:
        g (`torch.tensor`):
            gradient output tensor
        x (`torch.tensor`):
            input tensor


def save(fn: Callable[..., Dict[str, EventSetOrNode]], path: str, *args: Union[EventSetNode, EventSet, Schema], **kwargs: Union[EventSetNode, EventSet, Schema]) -> None:
    if not hasattr(fn, 'is_tp_compiled'):
        raise ValueError('Can only save a function that has been compiled with `@tp.compile`.')
    merged_kwargs = _kwargs_from_args_and_kwargs(list(inspect.signature(fn).parameters.keys()), args, kwargs)
    node_kwargs = {k: _process_fn_input(v) for k, v in merged_kwargs.items()}
    outputs = fn(**node_kwargs)
    _check_fn_outputs(outputs)
    save_graph(inputs=node_kwargs, outputs=outputs, path=path)

Saves a compiled Temporian function to a file.

    The saved function must only take
    [`EventSetOrNodes`][temporian.types.EventSetOrNode] as arguments, return a
    dictionary of names to EventSetOrNodes, and be decorated with
    [`@tp.compile`][temporian.compile].

    Temporian saves the graph built between the function's input and output
    nodes, not the function itself. Any arbitrary code that
    is executed in the function will not be ran when loading it back up and
    executing it.

    If you need to save a function that additionally takes other types of
    arguments, try using `functools.partial` to create a new function that takes
    only EventSetNodes, and save that instead. Note that the partial function
    needs to be compiled too, with `tp.compile(partial(...))`.

Args:
        fn: The function to save.
        path: The path to save the function to.
        args: Positional arguments to pass to the function to trace it. The
            arguments can be either EventSets, EventSetNodes, or raw Schemas. In
            all cases, the values will be converted to EventSetNodes before
            being passed to the function to trace it.
        kwargs: Keyword arguments to pass to the function to trace it. Same
            restrictions as for `args`.

Raises:
        ValueError: If the received function is not compiled.
        ValueError: If any of the received inputs is not of the specified types.
        ValueError: If the function doesn't return the specified type.


def MaxBipartiteMatching(self, graph):
    self.g = nx.Graph(graph)
    self.left = set((n for n, d in self.g.nodes(data=True) if not d['bipartite']))
    self.right = set(self.g) - self.left
    self.num_matched = 0
    self.s = set()
    self.t = set()
    self.matches = {}
    self.slack = {}
    self.slackx = {}
    self.prev = {}
    self.labels = {}
    for x in self.left:
        self.labels[x] = max([val['weight'] for val in self.g[x].values()])
    for y in self.right:
        self.labels[y] = 0
    while self.num_matched != len(self.left):
        self._Augment()
    ret = {}
    for k in self.left:
        ret[k] = self.matches[k]
    return ret

Find a maximum matching for a bipartite graph.

    This is O(n^3) implementation of the Hungarian method for complete bipartite
    matching problems.

Args:
        graph: A networkx graph object, assumed to be bipartite.

Returns:
        A dictionary keyed on node names in left to node names in right.


def __init__(self, stddev, seed=None, **kwargs):
    super().__init__(**kwargs)
    if not 0 <= stddev <= 1:
        raise ValueError(f'Invalid value received for argument `stddev`. Expected a float value between 0 and 1. Received: stddev={stddev}')
    self.stddev = stddev
    self.seed = seed
    if stddev > 0:
        self.seed_generator = backend.random.SeedGenerator(seed)
    self.supports_masking = True
    self._build_at_init()

Apply additive zero-centered Gaussian noise.

    This is useful to mitigate overfitting
    (you could see it as a form of random data augmentation).
    Gaussian Noise (GS) is a natural choice as corruption process
    for real valued inputs.

    As it is a regularization layer, it is only active at training time.

Args:
        stddev: Float, standard deviation of the noise distribution.
        seed: Integer, optional random seed to enable deterministic behavior.

    Call arguments:
        inputs: Input tensor (of any rank).
        training: Python boolean indicating whether the layer should behave in
            training mode (adding noise) or in inference mode (doing nothing).


def _kl_dirichlet_dirichlet(d1, d2, name=None):
    with ops.name_scope(name, 'kl_dirichlet_dirichlet', values=[d1.concentration, d2.concentration]):
        digamma_sum_d1 = math_ops.digamma(math_ops.reduce_sum(d1.concentration, axis=-1, keepdims=True))
        digamma_diff = math_ops.digamma(d1.concentration) - digamma_sum_d1
        concentration_diff = d1.concentration - d2.concentration
        return math_ops.reduce_sum(concentration_diff * digamma_diff, axis=-1) - special_math_ops.lbeta(d1.concentration) + special_math_ops.lbeta(d2.concentration)

Batchwise KL divergence KL(d1 || d2) with d1 and d2 Dirichlet.

Args:
        d1: instance of a Dirichlet distribution object.
        d2: instance of a Dirichlet distribution object.
        name: (optional) Name to use for created operations.
        default is "kl_dirichlet_dirichlet".

Returns:
        Batchwise KL(d1 || d2)


def _show_ops_in_metagraph(saved_model_dir, tag_set):
    meta_graph_def = saved_model_utils.get_meta_graph_def(saved_model_dir, tag_set)
    _show_ops_in_metagraph_mgd(meta_graph_def)

Prints the ops in the MetaGraph.

    Prints all the ops used in the MetaGraphDef indicated by the tag_set stored in
    SavedModel directory.

Args:
        saved_model_dir: Directory containing the SavedModel to inspect.
        tag_set: Group of tag(s) of the MetaGraphDef in string format, separated by
        ','. For tag-set contains multiple tags, all tags must be passed in.


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None, encoder_hidden_states: Optional[torch.FloatTensor]=None, encoder_attention_mask: Optional[torch.FloatTensor]=None, head_mask: Optional[torch.Tensor]=None, cross_attn_head_mask: Optional[torch.Tensor]=None, past_key_values: Optional[Tuple[Tuple[torch.Tensor]]]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, cache_position: Optional[torch.LongTensor]=None) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    outputs = self.model.decoder(input_ids=input_ids, attention_mask=attention_mask, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, head_mask=head_mask, cross_attn_head_mask=cross_attn_head_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position)
    logits = self.lm_head(outputs[0])
    loss = None
    if labels is not None:
        loss_fct = CrossEntropyLoss()
        loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
    if not return_dict:
        output = (logits,) + outputs[1:]
        return (loss,) + output if loss is not None else output
    return CausalLMOutputWithCrossAttentions(loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, cross_attentions=outputs.cross_attentions)

cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
    Mask to nullify selected heads of the cross-attention modules. Mask values selected in `[0, 1]`:

    - 1 indicates the head is **not masked**,
    - 0 indicates the head is **masked**.
    labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
    Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
    config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
    (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

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

    >>> tokenizer = AutoTokenizer.from_pretrained("google/bigbird-pegasus-large-arxiv")
    >>> model = BigBirdPegasusForCausalLM.from_pretrained(
    ...     "google/bigbird-pegasus-large-arxiv", add_cross_attention=False
    ... )
    >>> assert model.config.is_decoder, f"{model.__class__} has to be configured as a decoder."
    >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
    >>> outputs = model(**inputs)

    >>> logits = outputs.logits
    ```


def __init__(self, value_type: typing.Union[typing.Type[typing.Any], typing.Tuple[typing.Type[typing.Any], ...]], default: typing.Any=MISSING_VALUE, is_noneable: bool=False, frozen: bool=False):
    super().__init__(value_type, default, is_noneable=is_noneable, frozen=frozen)

Constructor.

Args:
        value_type: Acceptable value type(s).
        default: Default value.
        is_noneable: If True, None is acceptable.
        frozen: If True, values other than the default value is not accceptable.


def _boundaries_to_sizes(a, boundaries, axis):
    if axis >= len(a.shape):
        raise ValueError('axis %s is out of bound for shape %s' % (axis, a.shape))
    total_size = a.shape[axis]
    sizes = []
    sizes_sum = 0
    prev = 0
    for i, b in enumerate(boundaries):
        size = b - prev
        if size < 0:
            raise ValueError('The %s-th boundary %s is smaller than the previous boundary %s' % (i, b, prev))
        size = builtins.min(size, builtins.max(0, total_size - sizes_sum))
        sizes.append(size)
        sizes_sum += size
        prev = b
    sizes.append(builtins.max(0, total_size - sizes_sum))
    return sizes

Converting boundaries of splits to sizes of splits.

Args:
        a: the array to be split.
        boundaries: the boundaries, as in np.split.
        axis: the axis along which to split.

Returns:
        A list of sizes of the splits, as in tf.split.


def run_stages(self, stage_context: translations.TransformContext, stages: List[translations.Stage]) -> 'RunnerResult':
    worker_handler_manager = WorkerHandlerManager(stage_context.components.environments, self._provision_info)
    pipeline_metrics = MetricsContainer('')
    pipeline_metrics.get_counter(MetricName(str(type(self)), self.NUM_FUSED_STAGES_COUNTER, urn='internal:' + self.NUM_FUSED_STAGES_COUNTER)).update(len(stages))
    monitoring_infos_by_stage: MutableMapping[str, Iterable['metrics_pb2.MonitoringInfo']] = {}
    runner_execution_context = execution.FnApiRunnerExecutionContext(stages, worker_handler_manager, stage_context.components, stage_context.safe_coders, stage_context.data_channel_coders, self._num_workers, split_managers=self._split_managers)
    try:
        with self.maybe_profile():
            runner_execution_context.setup()
            bundle_counter = 0
            while len(runner_execution_context.queues.ready_inputs) > 0:
                _LOGGER.debug('Remaining ready bundles: %s\n\tWatermark pending bundles: %s\n\tTime pending bundles: %s', len(runner_execution_context.queues.ready_inputs), len(runner_execution_context.queues.watermark_pending_inputs), len(runner_execution_context.queues.time_pending_inputs))
                consuming_stage_name, bundle_input = runner_execution_context.queues.ready_inputs.deque()
                stage = runner_execution_context.stages[consuming_stage_name]
                bundle_context_manager = runner_execution_context.bundle_manager_for(stage)
                _BUNDLE_LOGGER.debug('Running bundle for stage %s\n\tExpected outputs: %s timers: %s', bundle_context_manager.stage.name, bundle_context_manager.stage_data_outputs, bundle_context_manager.stage_timer_outputs)
                assert consuming_stage_name == bundle_context_manager.stage.name
                bundle_counter += 1
                bundle_results = self._execute_bundle(runner_execution_context, bundle_context_manager, bundle_input)
                if consuming_stage_name in monitoring_infos_by_stage:
                    monitoring_infos_by_stage[consuming_stage_name] = consolidate_monitoring_infos(itertools.chain(bundle_results.process_bundle.monitoring_infos, monitoring_infos_by_stage[consuming_stage_name]))
                else:
                    assert isinstance(bundle_results.process_bundle.monitoring_infos, Iterable)
                    monitoring_infos_by_stage[consuming_stage_name] = bundle_results.process_bundle.monitoring_infos
                if '' not in monitoring_infos_by_stage:
                    monitoring_infos_by_stage[''] = list(pipeline_metrics.to_runner_api_monitoring_infos('').values())
                else:
                    monitoring_infos_by_stage[''] = consolidate_monitoring_infos(itertools.chain(pipeline_metrics.to_runner_api_monitoring_infos('').values(), monitoring_infos_by_stage['']))
                if len(runner_execution_context.queues.ready_inputs) == 0:
                    self._schedule_ready_bundles(runner_execution_context)
        assert len(runner_execution_context.queues.ready_inputs) == 0, 'A total of %d ready bundles did not execute.' % len(runner_execution_context.queues.ready_inputs)
        assert len(runner_execution_context.queues.watermark_pending_inputs) == 0, 'A total of %d watermark-pending bundles did not execute.' % len(runner_execution_context.queues.watermark_pending_inputs)
        assert len(runner_execution_context.queues.time_pending_inputs) == 0, 'A total of %d time-pending bundles did not execute.' % len(runner_execution_context.queues.time_pending_inputs)
    finally:
        worker_handler_manager.close_all()
    return RunnerResult(runner.PipelineState.DONE, monitoring_infos_by_stage)

Run a list of topologically-sorted stages in batch mode.

Args:
        stage_context (translations.TransformContext)
        stages (list[fn_api_runner.translations.Stage])


def offsets_in_rows(self):
    return gen_ragged_math_ops.ragged_range(starts=constant_op.constant(0, self.dtype), limits=self.row_lengths(), deltas=constant_op.constant(1, self.dtype)).rt_dense_values

Return the offset of each value.

    RowPartition takes an array x and converts it into sublists.
    offsets[i] is the index of x[i] in its sublist.
    Given a shape, such as:
    [*,*,*],[*,*],[],[*,*]
    This returns:
    0,1,2,0,1,0,1

Returns:
        an offset for every value.


def create_token_type_ids_from_sequences(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None) -> List[int]:
    sep = [self.sep_token_id]
    cls = [self.cls_token_id]
    if token_ids_1 is None:
        return len(cls) * [self.cls_token_type_id] + len(token_ids_0 + sep) * [0]
    return len(cls) * [self.cls_token_type_id] + len(token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]

Create a mask from the two sequences passed to be used in a sequence-pair classification task. A Funnel
    Transformer sequence pair mask has the following format:

    ```
    2 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
    | first sequence    | second sequence |
    ```

    If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).

Args:
        token_ids_0 (`List[int]`):
            List of IDs.
        token_ids_1 (`List[int]`, *optional*):
            Optional second list of IDs for sequence pairs.

Returns:
        `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).


def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.LongTensor]=None, past_key_value: Optional[Tuple[torch.Tensor]]=None, output_attentions: Optional[bool]=False, output_router_logits: Optional[bool]=False, use_cache: Optional[bool]=False, cache_position: Optional[torch.LongTensor]=None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]]=None, **kwargs) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
    residual = hidden_states
    hidden_states = self.input_layernorm(hidden_states)
    hidden_states, self_attn_weights, present_key_value = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings)
    hidden_states = residual + hidden_states
    residual = hidden_states
    hidden_states = self.post_attention_layernorm(hidden_states)
    hidden_states, router_logits = self.block_sparse_moe(hidden_states)
    hidden_states = residual + hidden_states
    outputs = (hidden_states,)
    if output_attentions:
        outputs += (self_attn_weights,)
    if use_cache:
        outputs += (present_key_value,)
    if output_router_logits:
        outputs += (router_logits,)
    return outputs

Args:
    hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
    attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
        `(batch, sequence_length)` where padding elements are indicated by 0.
    past_key_value (`Tuple(torch.FloatTensor)`, *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.
    output_router_logits (`bool`, *optional*):
        Whether or not to return the logits of all the routers. They are useful for computing the router loss, and
        should not be returned during inference.
    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`).
    cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
        Indices depicting the position of the input sequence tokens in the sequence.
    kwargs (`dict`, *optional*):
        Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
        into the model


def _experimental_default_mesh(self, mesh: layout_lib.Mesh):
    previous_default = self._current_default_mesh
    self._register_mesh(mesh)
    _pywrap_dtensor_device.ExperimentalSetDefaultMesh(self._device_info, mesh.to_string().encode('utf-8'))
    self._current_default_mesh = mesh
    yield
    _pywrap_dtensor_device.ExperimentalClearDefaultMesh(self._device_info)
    if previous_default:
        _pywrap_dtensor_device.ExperimentalSetDefaultMesh(self._device_info, previous_default.to_string().encode('utf-8'))
    self._current_default_mesh = previous_default

Sets a default mesh for all ops in the scope.

    Note: This is an internal helper method, which is not user facing api.

    Useful for requesting a specific mesh for ops which would have no inferred
    layout, e.g. tf.zeros.

Args:
        mesh: A Mesh to be used for ops without Mesh.

Yields:
        Nothing.


def lookup(self, obj):
    for registered in self._registry:
        if isinstance(obj, registered):
            return self._registry[registered]
    raise LookupError(f'{type(obj)} has not been registered.')

Looks up 'obj'.

Args:
        obj: The object to lookup within the registry.

Returns:
        Value for 'obj' in the registry if found.

Raises:
        LookupError: if 'obj' has not been registered.


def recipe_cm_user_editor(config, recipe_name):
    drive(config, {'auth': 'user', 'hour': [], 'copy': {'source': 'https://docs.google.com/spreadsheets/d/1Mw4kDJfaWVloyjSayJSkgE8i28Svoj1756fyQtIpmRE/', 'destination': recipe_name}})

A tool for rapidly bulk editing Campaign Manager profiles, roles, and sub
    accounts.

Args:
        recipe_name (string) - Name of document to deploy to.


def __add__(self, other):
    try:
        other = as_dimension(other)
    except (TypeError, ValueError):
        return NotImplemented
    if self._value is None or other.value is None:
        return Dimension(None)
    else:
        return Dimension(self._value + other.value)

Returns the sum of `self` and `other`.

    Dimensions are summed as follows:

    ```python
    tf.compat.v1.Dimension(m)    + tf.compat.v1.Dimension(n)     ==
    tf.compat.v1.Dimension(m + n)
    tf.compat.v1.Dimension(m)    + tf.compat.v1.Dimension(None)  # equiv. to
    tf.compat.v1.Dimension(None)
    tf.compat.v1.Dimension(None) + tf.compat.v1.Dimension(n)     # equiv. to
    tf.compat.v1.Dimension(None)
    tf.compat.v1.Dimension(None) + tf.compat.v1.Dimension(None)  # equiv. to
    tf.compat.v1.Dimension(None)
    ```

Args:
        other: Another Dimension, or a value accepted by `as_dimension`.

Returns:
        A Dimension whose value is the sum of `self` and `other`.


def in_range(self, ver, req):
    if req.exclude is not None:
        for v in ver:
            if v in req.exclude:
                return False
    include_checked = False
    if req.include is not None:
        for v in ver:
            if v in req.include:
                return True
        include_checked = True
    if req.range != [None, None]:
        min_v = req.range[0]
        max_v = req.range[1]
        ver = ver[0]
        lg = _compare_versions(min_v, ver)['larger']
        sm = _compare_versions(ver, max_v)['smaller']
        if lg in [ver, 'equal'] and sm in [ver, 'equal', 'inf']:
            return True
        else:
            err_msg = '[Error] Version is outside of supported range. '
            err_msg += '(config = %s, ' % str(req.config)
            err_msg += 'version = %s, ' % str(ver)
            err_msg += 'supported range = %s)' % str(req.range)
            logging.warning(err_msg)
            self.warning_msg.append(err_msg)
            return False
    else:
        err_msg = ''
        if include_checked:
            err_msg = '[Error] Version is outside of supported range. '
        else:
            err_msg = '[Error] Missing specification. '
        err_msg += '(config = %s, ' % str(req.config)
        err_msg += 'version = %s, ' % str(ver)
        err_msg += 'supported range = %s)' % str(req.range)
        logging.warning(err_msg)
        self.warning_msg.append(err_msg)
        return False

Checks if a version satisfies a version and/or compatibility requirement.

Args:
        ver: List whose first item is a config version that needs to be checked
        for support status and version compatibility.
            e.g. ver = [`1.0`]
        req: `_Reqs` class instance that represents a configuration version and
            compatibility specifications.

Returns:
        Boolean output of checking if version `ver` meets the requirement
        stored in `req` (or a `_Reqs` requirements class instance).


def write_line(self, values):
    if self.inner_table is None:
        self.inner_table = [list(values.keys()), list(values.values())]
    else:
        columns = self.inner_table[0]
        for key in values.keys():
            if key not in columns:
                columns.append(key)
        self.inner_table[0] = columns
        if len(self.inner_table) > 1:
            last_values = self.inner_table[-1]
            first_column = self.inner_table[0][0]
            if last_values[0] != values[first_column]:
                self.inner_table.append([values[c] if c in values else 'No Log' for c in columns])
            else:
                new_values = values
                for c in columns:
                    if c not in new_values.keys():
                        new_values[c] = last_values[columns.index(c)]
                self.inner_table[-1] = [new_values[c] for c in columns]
        else:
            self.inner_table.append([values[c] for c in columns])

Write the values in the inner table.

Args:
        values (`Dict[str, float]`): The values to display.


def __init__(self, data=None, _KEY=None, _ATTRS=None):
    if self.__class__ is MapEntry:
        raise TypeError('MapEntry is an abstract class.')
    if data is None:
        return
    else:
        for key in data:
            setattr(self, key, data[key])
    self.log = logging.getLogger(__name__)

This is an abstract class.

Args:
        data:  An optional dict of attribute, value pairs to populate with.

Raises:
        TypeError:  Bad argument, or attempt to instantiate abstract class.


def _set_device(self, device) -> None:
    self._set_device_from_string(compat.as_str(_device_string(device)))

Set the device of this operation.

Args:
        device: string or device..  The device to set.


def dummy_inputs(self) -> Dict[str, tf.Tensor]:
    return {self.main_input_name: tf.random.uniform([1, self.config.num_mel_bins, self.config.max_source_positions * 2 - 1], dtype=tf.float32), 'decoder_input_ids': tf.constant([[1, 3]], dtype=tf.int32)}

Dummy inputs to build the network.

Returns:
        `Dict[str, tf.Tensor]`: The dummy inputs.


def device(self) -> torch.device:
    if self._rot_mats is not None:
        return self._rot_mats.device
    elif self._quats is not None:
        return self._quats.device
    else:
        raise ValueError('Both rotations are None')

The device of the underlying rotation

Returns:
        The device of the underlying rotation


def get_range(self, request, start, end):
    file_ = self.get_file(request.bucket, request.object)
    if start < 0 or end <= start:
        return file_.contents
    return file_.contents[start:end]

Retrieves an object.

Args:
        request: (GetRequest) request

Returns:
        (bytes) The response message.


Called at the beginning of training.

    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 _make_output_composite_tensors_match(op_type, branch_graphs):
    assert branch_graphs
    branch_outputs = [g.structured_outputs for g in branch_graphs]
    outputs_per_branch = list((len(outs) for outs in branch_outputs))
    assert len(set(outputs_per_branch)) == 1, outputs_per_branch
    for output_idx, branch_outs in enumerate(zip(*branch_outputs)):
        if len(set((type(out) for out in branch_outs))) == 1:
            continue
        if not any((isinstance(out, indexed_slices.IndexedSlices) for out in branch_outs)):
            continue
        for branch_idx, branch_out in enumerate(branch_outs):
            if isinstance(branch_out, indexed_slices.IndexedSlices):
                continue
            elif isinstance(branch_out, tensor_lib.Tensor):
                with branch_graphs[branch_idx].as_default():
                    branch_outputs[branch_idx][output_idx] = math_ops._as_indexed_slices(branch_out)
            else:
                raise TypeError('Cannot reconcile {op_name} {output_idx}-th outputs:\n  outputs from all branches: {outputs}'.format(op_name='tf.cond' if op_type == _COND else 'tf.switch_case', output_idx=output_idx, outputs=branch_outs))
    for branch_graph, branch_outs in zip(branch_graphs, branch_outputs):
        branch_graph.structured_outputs = branch_outs
        branch_graph.outputs = [t for t in func_graph_module.flatten(branch_outs) if t is not None]

Modifies each branch_graph's outputs to have the same output signature.

    Currently the only transformation implemented is turning a Tensor into an
    equivalent IndexedSlices if the other branch returns an IndexedSlices.
    Updates branch_graph.{outputs,structured_outputs} for each branch_graph in
    branch_graphs.

Args:
        op_type: _COND or _CASE
        branch_graphs: `list` of `FuncGraph`

Raises:
        TypeError: if a set of outputs cannot be rewritten.


def _maybe_broadcast_to_outputs(self, outputs, objects):
    if not self._should_broadcast(objects):
        return objects
    should_copy_objects = len(nest.flatten(outputs)) > 1

    def _broadcast_fn():
        if should_copy_objects:
            return nest.map_structure(self._copy_object, objects)
        return objects
    return nest.map_structure(lambda _: _broadcast_fn(), outputs)

Determines if losses / metrics should be applied to all outputs.

    NOTE: This method should only be called for Metrics / Losses, not for
    y_true / sample_weight.

Args:
        outputs: Model predictions.
        objects: Arbitrary nested structure (e.g. of losses or metrics)

Returns:
        Arbitrary nested structure of objects, maybe copied to each output.

    Applies a Loss / Metric to all outputs.


def weights(self):
    return self._dedup_weights(self._undeduplicated_weights)

Returns the list of all layer variables/weights.

    Note: This will not track the weights of nested `tf.Modules` that are not
    themselves Keras layers.

Returns:
        A list of variables.


def _replace_tensors_for_gradient(x, grad):
    if not isinstance(x, composite_tensor.CompositeTensor):
        return grad
    if not isinstance(x, CompositeTensorGradientProtocol):
        raise ValueError(f'Type {type(x).__name__} is not supported as a gradient source.')
    composite_gradient = x.__composite_gradient__
    x_components = composite_gradient.get_gradient_components(x)
    if x_components is x:
        grad_components = grad
    else:
        grad_components = nest.map_structure_up_to(x_components, _replace_tensors_for_gradient, x_components, grad)
    if grad_components is None:
        return None
    return composite_gradient.replace_gradient_components(x, grad_components)

Replaces the tensors in `x` that should be differentiated with `grad`.

Args:
        x: A `Tensor` or `CompositeTensor`.
        grad: A nested structure of `Tensor`, with the same structure as the value
        returned by `_get_tensors_for_gradient(x)`.

Returns:
        A `Tensor` or `CompositeTensor`.


def get_image_size_fit_to_canvas(image_height: int, image_width: int, canvas_height: int, canvas_width: int, tile_size: int) -> Tuple[int, int]:
    target_width = np.clip(image_width, tile_size, canvas_width)
    target_height = np.clip(image_height, tile_size, canvas_height)
    scale_h = target_height / image_height
    scale_w = target_width / image_width
    if scale_w < scale_h:
        new_width = target_width
        new_height = min(math.floor(image_height * scale_w) or 1, target_height)
    else:
        new_height = target_height
        new_width = min(math.floor(image_width * scale_h) or 1, target_width)
    return (new_height, new_width)

Calculates the new size of an image to fit within a canvas while maintaining aspect ratio.

    This function calculates the optimal size for an image to fit within a canvas defined by
    canvas_height and canvas_width, while ensuring that the image dimensions are not smaller than
    tile_size. If the image is larger than the canvas, the returned size will fit within the canvas.
    If the image already fits within the canvas, the size remains unchanged.
    The aspect ratio of the original image is preserved as much as possible.

Args:
        image_height (`int`):
            The height of the original image.
        image_width (`int`):
            The width of the original image.
        canvas_height (`int`):
            The height of the canvas.
        canvas_width (`int`):
            The width of the canvas.
        tile_size (`int`):
            The tile size.

Returns:
        `Tuple[int, int]`: A tuple containing the new height and width of the image.


def sheets_tab_range(sheet_tab, sheet_range):
    if sheet_range:
        return '%s!%s' % (sheet_tab, sheet_range)
    else:
        return sheet_tab

Helper for creating range format.

Args:
        sheet_tab - name of tab in sheet
        sheet_range - A1 notation

Returns:
        String containing full sheet range specification.


def _CheckUnresolved(cls):
    with cls._unresolved_subjects_lock:
        if cls._unresolved_subjects:
            msg = ['The following assertions were unresolved. Perhaps you called "AssertThat(thing.IsEmpty())" instead of "AssertThat(thing).IsEmpty()".']
            for u in sorted(cls._unresolved_subjects):
                msg.append('    * {0}'.format(u))
            raise UnresolvedAssertionError('\n'.join(msg))

Ensures that all created subjects were eventually resolved.

    A subject is considered resolved what at least one proposition has been
    executed on it. An unresolved or dangling assertion is almost certainly a
    test author error.

Raises:
        UnresolvedAssertionError: if any subjects remain unresolved at the time of
        this function call.


def _apply_threshold_to_predictions(self, result: AnomalyResult) -> AnomalyResult:
    predictions = [dataclasses.replace(p, label=self._threshold_fn.apply(p.score), threshold=self._threshold_fn.threshold) for p in result.predictions]
    return dataclasses.replace(result, predictions=predictions)

Updates the prediction labels in an AnomalyResult using the ThresholdFn.

Args:
        result (AnomalyResult): The input `AnomalyResult` containing anomaly
        scores.

Returns:
        AnomalyResult: A new `AnomalyResult` with updated prediction labels
        and threshold values.


def set_iterator_element_layouts(self, iterator_resource_dtensor, layouts: List[layout_lib.Layout]):
    _pywrap_dtensor_device.SetIteratorElementLayouts(context.context()._handle, iterator_resource_dtensor, [layout.to_string() for layout in layouts], self._device_info)

Sets the element layouts on an iterator resource tensor.

Args:
        iterator_resource_dtensor: a DTensor created by packing the individiual
        iterator resource tensors.
        layouts: the flattened list of layouts to be applied to the elements
        emitted by the iterator resource DTensor.


def _add_file_argument(parser, types, args, custom_kwargs=None):
    custom_kwargs = custom_kwargs or {}
    arg = args[-1]
    dest = custom_kwargs.get('dest', arg.lstrip(_ARG_PREFIX).replace('-', '_'))
    kwargs = {'type': types.get(dest), 'action': 'store', 'default': config.ITEMS[dest].default, 'help': config.ITEMS[dest].comment}
    kwargs.update(custom_kwargs)
    if kwargs['type'] is None:
        del kwargs['type']
    if arg.startswith(_ARG_PREFIX):
        kwargs['dest'] = dest
    elif 'type' in kwargs:
        kwargs['default'] = kwargs['type'](kwargs['default'])
    parser.add_argument(*args, **kwargs)

Add a file-configurable option to the parser.

Args:
        parser: The parser.
        types: A map from option destination to type.
        args: The option's name(s). Either a 2-tuple of (short_arg, arg) or a
        1-tuple of (arg,).
        custom_kwargs: The option's custom kwargs.


def crop_image_to_patches(self, images: 'torch.Tensor', min_patches: int, max_patches: int, use_thumbnail: bool=True, patch_size: Optional[Union[Tuple, int, dict]]=None, interpolation: Optional['F.InterpolationMode']=None):
    patch_size_height, patch_size_width = (patch_size.height, patch_size.width)
    original_height, original_width = images.shape[-2:]
    num_columns, num_rows = get_optimal_tiled_canvas((original_height, original_width), (patch_size_height, patch_size_width), min_patches, max_patches)
    target_width = patch_size_width * num_columns
    target_height = patch_size_height * num_rows
    num_blocks = num_columns * num_rows
    resized_image = self.resize(images, SizeDict(height=target_height, width=target_width), interpolation=interpolation)
    processed_images = []
    for i in range(num_blocks):
        column = i % num_columns
        row = i // num_columns
        box = (column * patch_size_width, row * patch_size_height, (column + 1) * patch_size_width, (row + 1) * patch_size_height)
        patch_image = resized_image[..., box[1]:box[3], box[0]:box[2]]
        processed_images.append(patch_image)
    if use_thumbnail and len(processed_images) != 1:
        thumbnail_img = self.resize(images, patch_size, interpolation=interpolation)
        processed_images.append(thumbnail_img)
    processed_images = torch.stack(processed_images, dim=0).transpose(0, 1).contiguous()
    return processed_images

Crop the images to patches and return a list of cropped images.
    The number of patches and their grid arrangement are determined by the original image size,
    the target patch size and the minimum and maximum number of patches.
    The aspect ratio of the patches grid is chosen to be the closest to the original image aspect ratio.

Args:
        images (`torch.Tensor`):
            The images to be cropped.
        min_patches (`int`):
            The minimum number of patches to be extracted from the image.
        max_patches (`int`):
            The maximum number of patches to be extracted from the image.
        use_thumbnail (`bool`, *optional*, defaults to `True`):
            Whether to add a thumbnail image to the list of cropped patches.
        patch_size (`int`, `Tuple[int, int]`, `dict`, *optional*):
            The size of the output patches.
            The format of the image data. If `None`, the format is inferred from the input image.

Returns:
        List[`PIL.Image.Image`] or List[np.ndarray]: The list of cropped images.


def copy_file(src_file: str, dst_dir: str, strip: str=None, dest_file: str=None) -> None:
    dest = dest_file if dest_file else src_file
    if dest.startswith('bazel-out'):
        dest = dest[dest.index('bin') + 4:]
    if strip:
        dest = dest.removeprefix(strip)
    dest_dir_path = os.path.join(dst_dir, os.path.dirname(dest))
    os.makedirs(dest_dir_path, exist_ok=True)
    shutil.copy(src_file, dest_dir_path)
    os.chmod(os.path.join(dst_dir, dest), 420)

Copy a file to the destination directory.

Args:
        src_file: file to be copied
        dst_dir: destination directory
        strip: prefix to strip before copying to destination
        dest_file: destanation file location if different from src_file


def _hertz_to_mel(frequencies_hertz, name=None):
    with ops.name_scope(name, 'hertz_to_mel', [frequencies_hertz]):
        frequencies_hertz = ops.convert_to_tensor(frequencies_hertz)
        return _MEL_HIGH_FREQUENCY_Q * math_ops.log(1.0 + frequencies_hertz / _MEL_BREAK_FREQUENCY_HERTZ)

Converts frequencies in `frequencies_hertz` in Hertz to the mel scale.

Args:
        frequencies_hertz: A `Tensor` of frequencies in Hertz.
        name: An optional name for the operation.

Returns:
        A `Tensor` of the same shape and type of `frequencies_hertz` containing
        frequencies in the mel scale.


def __init__(self, input_dataset, num_workers, index, num_replicas=None):
    self._input_dataset = input_dataset
    self._element_spec = input_dataset.element_spec
    variant_tensor = ged_ops.auto_shard_dataset(self._input_dataset._variant_tensor, num_workers=num_workers, index=index, auto_shard_policy=int(input_dataset.options().experimental_distribute.auto_shard_policy), num_replicas=num_replicas, **self._flat_structure)
    super(_AutoShardDataset, self).__init__(input_dataset, variant_tensor)

A `Dataset` that shards the `Dataset` automatically.

    This dataset takes in an existing dataset and tries to automatically figure
    out how to shard the dataset in a multi-worker scenario using graph rewrites.

    If the AutoShardPolicy is set to FILE, it walks up the dataset graph until
    it finds a reader dataset, then inserts a ShardDataset op before that node
    so that each worker only sees some files.

    If the AutoShardPolicy is set to DATA, it inserts a ShardDataset op at the
    end of the input pipeline, before any terminal PrefetchDataset if there is
    one. Additionally, if there is a RebatchDatasetV2 in the input pipeline, it
    is written to legacy RebatchDataset for correctness reasons, since
    RebatchDatasetV2 is incompatible with data sharding.

    If the AutoShardPolicy is set to AUTO, it tries to do file-based sharding.
    If it cannot find a reader dataset, it falls back to doing data-based
    sharding.

    If the AutoShardPolicy is set to OFF, it does nothing.

Attributes:
        num_workers: Total number of workers to shard this dataset across.
        index: The current worker index (out of the total number of workers) this
        dataset is for.
        num_replicas: The total number of replicas across all workers. This is used
        only when sharding by data (either DATA or AUTO) in order to rewrite
        RebatchDatasetV2 to RebatchDataset.

Raises:
        NotFoundError: If we cannot find a suitable reader dataset to begin
        automatically sharding the dataset.


def _add_name_scope_wrapper(func, api_signature):
    if 'name' not in api_signature.parameters:
        return func
    func_signature = tf_inspect.signature(func)
    func_argspec = tf_inspect.getargspec(func)
    if 'name' in func_signature.parameters or func_argspec.keywords is not None:
        return func
    name_index = list(api_signature.parameters).index('name')

    def wrapped_func(*args, **kwargs):
        if name_index < len(args):
            name = args[name_index]
            args = args[:name_index] + args[name_index + 1:]
        else:
            name = kwargs.pop('name', None)
        if name is None:
            return func(*args, **kwargs)
        else:
            with ops.name_scope(name):
                return func(*args, **kwargs)
    wrapped_func = tf_decorator.make_decorator(func, wrapped_func)
    wrapped_func.__signature__ = func_signature.replace(parameters=list(func_signature.parameters.values()) + [api_signature.parameters['name']])
    del wrapped_func._tf_decorator
    return wrapped_func

Wraps `func` to expect a "name" arg, and use it to call `ops.name_scope`.

    If `func` already expects a "name" arg, or if `api_signature` does not
    expect a "name" arg, then returns `func` as-is.

Args:
        func: The function to wrap.  Signature must match `api_signature` (except
        the "name" parameter may be missing.
        api_signature: The signature of the original API (used to find the index for
        the "name" parameter).

Returns:
        The wrapped function (or the original function if no wrapping is needed).


def hasValue(self) -> 'Builder':
    return self._to_builder(_evaluation.HasValueFunction(self.node.context, self.node, []))

The FHIRPath hasValue() function.

Returns:
        An expression that evaluates to True if the parent has a single value
        that is a primitive.


def reset_logical_devices(device_type, count):
    reset_context()
    devices = tf_config.list_physical_devices(device_type)
    if device_type.upper() not in ('CPU', 'GPU'):
        raise ValueError('resetting logical device for non-supported device type : %s' % device_type)
    if count < len(devices):
        devices = devices[:count]
        tf_config.set_visible_devices(devices, device_type=device_type.upper())
    for i, device in enumerate(devices):
        n = (i + 1) * count // len(devices) - i * count // len(devices)
        assert n > 0
        configs = []
        for ordinal in range(n):
            if device_type.upper() == 'GPU':
                dev_config = context.LogicalDeviceConfiguration(memory_limit=_DEFAULT_GPU_MEMORY_LIMIT, experimental_device_ordinal=ordinal)
            else:
                dev_config = context.LogicalDeviceConfiguration()
            configs.append(dev_config)
        tf_config.set_logical_device_configuration(device, configs)

Resets logical devices for CPU/GPU.

    Logical devices can only be instantiated once on a particular context. For
    now, context re-use is triggering some function duplication errors, so we
    reset the context on each call.

Args:
        device_type: The device_type to reset.
        count: numbers of virtual device to reset to.