def recipe_dcm_run(config, auth_read, account, report_id, report_name):
    dcm(config, {'auth': auth_read, 'report_run_only': True, 'report': {'account': account, 'report_id': report_id, 'name': report_name}})

Trigger a CM report run

Args:
  auth_read (authentication) - Credentials used for reading data.
  account (integer) - CM network id.
  report_id (integer) - CM report id, empty if using name.
  report_name (string) - CM report name, empty if using id instead.


def _rewrite_input_as_indexed_slices(body_grad_graph, grad_output_slices, forward_input, loop_vars):
    init_slices = _create_grad_indexed_slices_init(grad_output_slices, forward_input)
    with body_grad_graph.as_default():
        input_slices = indexed_slices.IndexedSlices(values=body_grad_graph.capture(init_slices.values, allowlisted=True), indices=body_grad_graph.capture(init_slices.indices, allowlisted=True), dense_shape=body_grad_graph.capture(init_slices.dense_shape, allowlisted=True))
        for t in _flatten(init_slices):
            captured_t = body_grad_graph.captures.pop(t)
            body_grad_graph.inputs.remove(captured_t)
        new_output_slices = _rewrite_grad_indexed_slices_output(grad_output_slices, input_slices)
    return _update_indexed_slices_param(body_grad_graph, loop_vars, init_slices, input_slices, new_output_slices, grad_output_slices)

Rewrites grad_output_slices's corresponding input to be an IndexedSlices.

This rewrite requires that forward_input was captured in the forward loop,
i.e. is not a user-specified loop variable. This is important because the
rewrite assumes that forward_input is passed through to its corresponding
output unchanged. This assumption is used in _rewrite_input_as_indexed_slices,
which depends on the exact gradient structure produced by the input's fanout.

This can yield a more efficient computation than using
_rewrite_output_as_tensor, since it preserves the IndexedSlices structure
instead of converting the IndexedSlices to a dense Tensor.

Args:
  body_grad_graph: _WhileBodyGradFuncGraph.
  grad_output_slices: IndexedSlices output of body_grad_graph.
  forward_input: the corresponding Tensor input to the forward loop.
  loop_vars: list of Tensors. The inputs to body_grad_graph.

Returns:
  The new loop_vars to pass to body_grad_graph.


def to_json(self, **kwargs):
    from keras.src.saving import serialization_lib
    model_config = serialization_lib.serialize_keras_object(self)
    return json.dumps(model_config, **kwargs)

Returns a JSON string containing the network configuration.

To load a network from a JSON save file, use
`keras.models.model_from_json(json_string, custom_objects={...})`.

Args:
    **kwargs: Additional keyword arguments to be passed to
        `json.dumps()`.

Returns:
    A JSON string.


def init_cache(self, batch_size, max_length, encoder_outputs):
    decoder_input_ids = jnp.ones((batch_size, max_length), dtype='i4')
    decoder_attention_mask = jnp.ones_like(decoder_input_ids)
    decoder_position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(decoder_input_ids).shape[-1]), decoder_input_ids.shape)

    def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):
        decoder_module = module._get_decoder_module()
        return decoder_module(decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs)
    init_variables = self.module.init(jax.random.PRNGKey(0), decoder_input_ids=decoder_input_ids, decoder_attention_mask=decoder_attention_mask, decoder_position_ids=decoder_position_ids, encoder_hidden_states=encoder_outputs[0], init_cache=True, method=_decoder_forward)
    return unfreeze(init_variables['cache'])

Args:
    batch_size (`int`):
        batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.
    max_length (`int`):
        maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized
        cache.
    encoder_outputs (`Union[FlaxBaseModelOutput, tuple(tuple(jnp.ndarray)]`):
        `encoder_outputs` consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*:
        `attentions`). `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*)
        is a sequence of hidden-states at the output of the last layer of the encoder. Used in the
        cross-attention of the decoder.


def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, position_embeddings: Optional[torch.Tensor]=None, output_attentions: bool=False, **kwargs):
    residual = hidden_states
    if self.normalize_before:
        hidden_states = self.self_attn_layer_norm(hidden_states)
    hidden_states, attn_weights = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask, position_embeddings=position_embeddings, output_attentions=output_attentions)
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    hidden_states = residual + hidden_states
    if not self.normalize_before:
        hidden_states = self.self_attn_layer_norm(hidden_states)
    if self.normalize_before:
        hidden_states = self.final_layer_norm(hidden_states)
    residual = hidden_states
    hidden_states = self.activation_fn(self.fc1(hidden_states))
    hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
    hidden_states = self.fc2(hidden_states)
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    hidden_states = residual + hidden_states
    if not self.normalize_before:
        hidden_states = self.final_layer_norm(hidden_states)
    if self.training:
        if torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any():
            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
    outputs = (hidden_states,)
    if output_attentions:
        outputs += (attn_weights,)
    return outputs

Args:
    hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
    attention_mask (`torch.FloatTensor`): attention mask of size
        `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative
        values.
    position_embeddings (`torch.FloatTensor`, *optional*):
        Object queries (also called content embeddings), to be added to the hidden 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.


class ProgramContext(object):

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

ProgramContext keeps track of converting function hierarchies.

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


def forward(self, input, tokens_per_expert):
    return sequential_experts_gemm(input, self.weight, tokens_per_expert.cpu())

Perform grouped matrix multiplication.

Args:
    input (`torch.Tensor`):
        Input tensor of shape (num_tokens, in_features).
    tokens_per_expert (`torch.Tensor`):
        Number of tokens assigned to each expert.

Returns:
    torch.Tensor: Output tensor of shape (num_tokens, out_features).


def from_dict(event_dict):
    return CallbackEvent(callback_id=event_dict['callbackId'], name=event_dict['name'], creation_time=event_dict['time'], data=event_dict['data'])

Creates a CallbackEvent object from a dictionary.

Args:
  event_dict: dict, a dictionary representing an event.

Returns:
  A CallbackEvent object.


class AriaTextConfig(LlamaConfig):
    model_type = 'aria_text'
    base_config_key = 'text_config'

    def __init__(self, intermediate_size: int=4096, moe_num_experts: int=8, moe_topk: int=2, moe_num_shared_experts: int=2, pad_token_id=2, **super_kwargs):
        super().__init__(pad_token_id=pad_token_id, **super_kwargs)
        self.intermediate_size = intermediate_size
        self.moe_num_experts = moe_num_experts
        self.moe_topk = moe_topk
        self.moe_num_shared_experts = moe_num_shared_experts

This class handles the configuration for the text component of the Aria model.
Instantiating a configuration with the defaults will yield a similar configuration to that of the model of the Aria
[rhymes-ai/Aria](https://huggingface.co/rhymes-ai/Aria) architecture.
This class extends the LlamaConfig to include additional parameters specific to the Mixture of Experts (MoE) architecture.

Args:
    vocab_size (`int`, *optional*, defaults to 32000):
        Vocabulary size of the LLaMA model. Defines the number of different tokens that can be represented by the
        `inputs_ids` passed when calling [`LlamaModel`]
    hidden_size (`int`, *optional*, defaults to 4096):
        Dimension of the hidden representations.
    intermediate_size (`int`, *optional*, defaults to 4096):
        The size of the MLP representations.
    num_hidden_layers (`int`, *optional*, defaults to 32):
        Number of hidden layers in the Transformer decoder.
    num_attention_heads (`int`, *optional*, defaults to 32):
        Number of attention heads for each attention layer in the Transformer decoder.
    num_key_value_heads (`int`, *optional*):
        This is the number of key_value heads that should be used to implement Grouped Query Attention. If
        `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
        `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
        converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
        by meanpooling all the original heads within that group. For more details, check out [this
        paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to
        `num_attention_heads`.
    hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
        The non-linear activation function (function or string) in the decoder.
    max_position_embeddings (`int`, *optional*, defaults to 2048):
        The maximum sequence length that this model might ever be used with. Llama 1 supports up to 2048 tokens,
        Llama 2 up to 4096, CodeLlama up to 16384.
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    rms_norm_eps (`float`, *optional*, defaults to 1e-06):
        The epsilon used by the rms normalization layers.
    use_cache (`bool`, *optional*, defaults to `True`):
        Whether or not the model should return the last key/values attentions (not used by all models). Only
        relevant if `config.is_decoder=True`.
    pad_token_id (`int`, *optional*, defaults to 2):
        Padding token id.
    bos_token_id (`int`, *optional*, defaults to 1):
        Beginning of stream token id.
    eos_token_id (`int`, *optional*, defaults to 2):
        End of stream token id.
    pretraining_tp (`int`, *optional*, defaults to 1):
        Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
        document](https://huggingface.co/docs/transformers/main/perf_train_gpu_many#tensor-parallelism) to
        understand more about it. This value is necessary to ensure exact reproducibility of the pretraining
        results. Please refer to [this issue](https://github.com/pytorch/pytorch/issues/76232).
    tie_word_embeddings (`bool`, *optional*, defaults to `False`):
        Whether to tie weight embeddings
    rope_theta (`float`, *optional*, defaults to 10000.0):
        The base period of the RoPE embeddings.
    rope_scaling (`Dict`, *optional*):
        Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
        and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
        accordingly.
        Expected contents:
            `rope_type` (`str`):
                The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
                'llama3'], with 'default' being the original RoPE implementation.
            `factor` (`float`, *optional*):
                Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
                most scaling types, a `factor` of x will enable the model to handle sequences of length x *
                original maximum pre-trained length.
            `original_max_position_embeddings` (`int`, *optional*):
                Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
                pretraining.
            `attention_factor` (`float`, *optional*):
                Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
                computation. If unspecified, it defaults to value recommended by the implementation, using the
                `factor` field to infer the suggested value.
            `beta_fast` (`float`, *optional*):
                Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
                ramp function. If unspecified, it defaults to 32.
            `beta_slow` (`float`, *optional*):
                Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
                ramp function. If unspecified, it defaults to 1.
            `short_factor` (`List[float]`, *optional*):
                Only used with 'longrope'. The scaling factor to be applied to short contexts (<
                `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
                size divided by the number of attention heads divided by 2
            `long_factor` (`List[float]`, *optional*):
                Only used with 'longrope'. The scaling factor to be applied to long contexts (<
                `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
                size divided by the number of attention heads divided by 2
            `low_freq_factor` (`float`, *optional*):
                Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
            `high_freq_factor` (`float`, *optional*):
                Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
    attention_bias (`bool`, *optional*, defaults to `False`):
        Whether to use a bias in the query, key, value and output projection layers during self-attention.
    attention_dropout (`float`, *optional*, defaults to 0.0):
        The dropout ratio for the attention probabilities.
    mlp_bias (`bool`, *optional*, defaults to `False`):
        Whether to use a bias in up_proj, down_proj and gate_proj layers in the MLP layers.
    head_dim (`int`, *optional*):
        The attention head dimension. If None, it will default to hidden_size // num_heads
    moe_num_experts (`int`, *optional*, defaults to 8):
        The number of experts in the MoE layer.
    moe_topk (`int`, *optional*, defaults to 2):
        The number of top experts to route to for each token.
    moe_num_shared_experts (`int`, *optional*, defaults to 2):
        The number of shared experts.


def call(self, input_ids: tf.Tensor | None=None, attention_mask: tf.Tensor | None=None, decoder_input_ids: tf.Tensor | None=None, decoder_attention_mask: tf.Tensor | None=None, decoder_position_ids: tf.Tensor | None=None, head_mask: tf.Tensor | None=None, decoder_head_mask: tf.Tensor | None=None, cross_attn_head_mask: tf.Tensor | None=None, encoder_outputs: Optional[Union[Tuple, TFBaseModelOutput]]=None, past_key_values: List[tf.Tensor] | None=None, inputs_embeds: tf.Tensor | None=None, decoder_inputs_embeds: tf.Tensor | None=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, labels: tf.Tensor | None=None, training: Optional[bool]=False) -> Union[Tuple[tf.Tensor], TFSeq2SeqLMOutput]:
    if labels is not None:
        labels = tf.where(labels == self.config.pad_token_id, tf.cast(tf.fill(shape_list(labels), -100), labels.dtype), labels)
        use_cache = False
        if decoder_input_ids is None and decoder_inputs_embeds is None:
            decoder_input_ids = shift_tokens_right(labels, self.config.pad_token_id, self.config.decoder_start_token_id)
    outputs = self.model(input_ids, attention_mask=attention_mask, decoder_input_ids=decoder_input_ids, encoder_outputs=encoder_outputs, decoder_attention_mask=decoder_attention_mask, decoder_position_ids=decoder_position_ids, head_mask=head_mask, decoder_head_mask=decoder_head_mask, cross_attn_head_mask=cross_attn_head_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, training=training)
    lm_logits = tf.matmul(outputs[0], self.model.shared.weights, transpose_b=True)
    lm_logits = self.bias_layer(lm_logits)
    masked_lm_loss = None if labels is None else self.hf_compute_loss(labels, lm_logits)
    if not return_dict:
        output = (lm_logits,) + outputs[1:]
        return (masked_lm_loss,) + output if masked_lm_loss is not None else output
    return TFSeq2SeqLMOutput(loss=masked_lm_loss, logits=lm_logits, past_key_values=outputs.past_key_values, decoder_hidden_states=outputs.decoder_hidden_states, decoder_attentions=outputs.decoder_attentions, cross_attentions=outputs.cross_attentions, encoder_last_hidden_state=outputs.encoder_last_hidden_state, encoder_hidden_states=outputs.encoder_hidden_states, encoder_attentions=outputs.encoder_attentions)

labels (`tf.tensor` 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]`.

Returns:


def switch(condition, then_expression, else_expression):
    if condition.dtype != dtypes_module.bool:
        condition = math_ops.cast(condition, 'bool')
    cond_ndim = ndim(condition)
    if not cond_ndim:
        if not callable(then_expression):

            def then_expression_fn():
                return then_expression
        else:
            then_expression_fn = then_expression
        if not callable(else_expression):

            def else_expression_fn():
                return else_expression
        else:
            else_expression_fn = else_expression
        x = cond.cond(condition, then_expression_fn, else_expression_fn)
    else:
        if callable(then_expression):
            then_expression = then_expression()
        if callable(else_expression):
            else_expression = else_expression()
        expr_ndim = ndim(then_expression)
        if cond_ndim > expr_ndim:
            raise ValueError('Rank of `condition` should be less than or equal to rank of `then_expression` and `else_expression`. ndim(condition)=' + str(cond_ndim) + ', ndim(then_expression)=' + str(expr_ndim))
        if cond_ndim > 1:
            ndim_diff = expr_ndim - cond_ndim
            cond_shape = array_ops.concat([array_ops.shape(condition), [1] * ndim_diff], axis=0)
            condition = array_ops.reshape(condition, cond_shape)
            expr_shape = array_ops.shape(then_expression)
            shape_diff = expr_shape - cond_shape
            tile_shape = array_ops.where_v2(shape_diff > 0, expr_shape, array_ops.ones_like(expr_shape))
            condition = array_ops.tile(condition, tile_shape)
        x = array_ops.where_v2(condition, then_expression, else_expression)
    return x

Switches between two operations depending on a scalar value.

Note that both `then_expression` and `else_expression`
should be symbolic tensors of the *same shape*.

Args:
    condition: tensor (`int` or `bool`).
    then_expression: either a tensor, or a callable that returns a tensor.
    else_expression: either a tensor, or a callable that returns a tensor.

Returns:
    The selected tensor.

Raises:
    ValueError: If rank of `condition` is greater than rank of expressions.


def copy_common_field(source_message: message.Message, target_message: message.Message, field_name: str):
    source_descriptor = source_message.DESCRIPTOR
    target_descriptor = target_message.DESCRIPTOR
    source_field = _field_descriptor_for_name(source_message, field_name)
    target_field = _field_descriptor_for_name(target_message, field_name)
    if source_field.type != target_field.type:
        raise ValueError(f'Field {field_name} differs in type between {source_descriptor.full_name} ({source_field.type}) and {target_descriptor.full_name} ({target_field.type}).')
    if field_is_repeated(source_field) != field_is_repeated(target_field):
        raise ValueError(f'Field {field_name} differs in size between {source_descriptor.full_name} and {target_descriptor.full_name}.')
    if field_is_set(source_message, source_field):
        source_value = get_value_at_field(source_message, source_field)
        set_value_at_field(target_message, target_field, source_value)

Copies field named field_name from source_message to target_message.

Args:
  source_message: The message to copy values from.
  target_message: The message to copy values to.
  field_name: The common field to examine/copy from source_message to
    target_message.

Raises:
  ValueError: <source> is not the same type as <target>. Unable to copy field
  <field>.
  ValueError: Field <field> is not present in both <source> and <target>.
  ValueError: Field <field> differs in type between <source> (<type>) and
  <target> (<type>).
  ValueError: Field <field> differs in size between <source> and <target>.


def loss_masks(self, masks_queries_logits: torch.Tensor, mask_labels: List[torch.Tensor], indices: Tuple[np.array], num_masks: int) -> Dict[str, torch.Tensor]:
    src_idx = self._get_predictions_permutation_indices(indices)
    tgt_idx = self._get_targets_permutation_indices(indices)
    pred_masks = masks_queries_logits[src_idx]
    target_masks, _ = self._pad_images_to_max_in_batch(mask_labels)
    target_masks = target_masks[tgt_idx]
    pred_masks = pred_masks[:, None]
    target_masks = target_masks[:, None]
    with torch.no_grad():
        point_coordinates = self.sample_points_using_uncertainty(pred_masks, lambda logits: self.calculate_uncertainty(logits), self.num_points, self.oversample_ratio, self.importance_sample_ratio)
        point_labels = sample_point(target_masks, point_coordinates, align_corners=False).squeeze(1)
    point_logits = sample_point(pred_masks, point_coordinates, align_corners=False).squeeze(1)
    losses = {'loss_mask': sigmoid_cross_entropy_loss(point_logits, point_labels, num_masks), 'loss_dice': dice_loss(point_logits, point_labels, num_masks)}
    del pred_masks
    del target_masks
    return losses

Compute the losses related to the masks using sigmoid_cross_entropy_loss and dice loss.

Args:
    masks_queries_logits (`torch.Tensor`):
        A tensor of shape `(batch_size, num_queries, height, width)`.
    mask_labels (`torch.Tensor`):
        List of mask labels of shape `(labels, height, width)`.
    indices (`Tuple[np.array])`:
        The indices computed by the Hungarian matcher.
    num_masks (`int)`:
        The number of masks, used for normalization.

Returns:
    losses (`Dict[str, Tensor]`): A dict of `torch.Tensor` containing two keys:
    - **loss_mask** -- The loss computed using sigmoid cross entropy loss on the predicted and ground truth.
      masks.
    - **loss_dice** -- The loss computed using dice loss on the predicted on the predicted and ground truth,
      masks.


def __getitem__(self, key: Any) -> 'ColumnExpressionBuilder':
    item = self._builder[key]
    if isinstance(item, expressions.Builder) and self._sealed:
        raise self._fhir_path_sealed_error(key)
    return ColumnExpressionBuilder._wrap_any(self, item)

Redirects to the expressions.Builder to get the item.

Args:
  key: the key of the item.

Returns:
  A ColumnExpressionBuilder, because the item got from the
  expressions.Builder is always the type of Builder.

Raises:
  AttributeError: if the FHIR path in this class is already sealed.
  TypeError: if getting the key from self._builder fails.


def normalize_element(element, element_signature=None):
    normalized_components = []
    if element_signature is None:
        components = nest.flatten(element)
        flattened_signature = [None] * len(components)
        pack_as = element
    else:
        flattened_signature = nest.flatten(element_signature)
        components = nest.flatten_up_to(element_signature, element)
        pack_as = element_signature
    with ops.name_scope('normalize_element'):
        for i, (t, spec) in enumerate(zip(components, flattened_signature)):
            try:
                if spec is None:
                    spec = type_spec_from_value(t, use_fallback=False)
            except TypeError:
                normalized_components.append(ops.convert_to_tensor(t, name='component_%d' % i))
            else:
                if spec.__class__.__name__ == 'DatasetSpec':
                    normalized_components.append(t)
                elif isinstance(spec, sparse_tensor.SparseTensorSpec):
                    normalized_components.append(sparse_tensor.SparseTensor.from_value(t))
                elif isinstance(spec, ragged_tensor.RaggedTensorSpec):
                    normalized_components.append(ragged_tensor.convert_to_tensor_or_ragged_tensor(t, name='component_%d' % i))
                elif isinstance(spec, tensor_array_ops.TensorArraySpec):
                    normalized_components.append(t)
                elif isinstance(spec, none_tensor.NoneTensorSpec):
                    normalized_components.append(none_tensor.NoneTensor())
                elif isinstance(spec, resource_variable_ops.VariableSpec):
                    normalized_components.append(ops.convert_to_tensor(t, name=f'component_{i}', dtype=spec.dtype))
                elif isinstance(t, composite_tensor.CompositeTensor):
                    normalized_components.append(t)
                else:
                    dtype = getattr(spec, 'dtype', None)
                    normalized_components.append(ops.convert_to_tensor(t, name='component_%d' % i, dtype=dtype))
    return nest.pack_sequence_as(pack_as, normalized_components)

Normalizes a nested structure of element components.

* Components matching `SparseTensorSpec` are converted to `SparseTensor`.
* Components matching `RaggedTensorSpec` are converted to `RaggedTensor`.
* Components matching `VariableSpec` are converted to `Tensor`.
* Components matching `DatasetSpec` or `TensorArraySpec` are passed through.
* `CompositeTensor` components are passed through.
* All other components are converted to `Tensor`.

Args:
  element: A nested structure of individual components.
  element_signature: (Optional.) A nested structure of `tf.DType` objects
    corresponding to each component of `element`. If specified, it will be
    used to set the exact type of output tensor when converting input
    components which are not tensors themselves (e.g. numpy arrays, native
    python types, etc.)

Returns:
  A nested structure of `Tensor`, `Variable`, `Dataset`, `SparseTensor`,
  `RaggedTensor`, or `TensorArray` objects.


class XLMPoolerEndLogits(nn.Module):

    def __init__(self, config: XLMConfig):
        super().__init__()
        self.dense_0 = nn.Linear(config.hidden_size * 2, config.hidden_size)
        self.activation = nn.Tanh()
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dense_1 = nn.Linear(config.hidden_size, 1)

    def forward(self, hidden_states: torch.FloatTensor, start_states: Optional[torch.FloatTensor]=None, start_positions: Optional[torch.LongTensor]=None, p_mask: Optional[torch.FloatTensor]=None) -> torch.FloatTensor:
        """
        Args:
            hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`):
                The final hidden states of the model.
            start_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`, *optional*):
                The hidden states of the first tokens for the labeled span.
            start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
                The position of the first token for the labeled span.
            p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*):
                Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token
                should be masked.

        <Tip>

        One of `start_states` or `start_positions` should be not `None`. If both are set, `start_positions` overrides
        `start_states`.

        </Tip>

        Returns:
            `torch.FloatTensor`: The end logits for SQuAD.
        """
        assert start_states is not None or start_positions is not None, 'One of start_states, start_positions should be not None'
        if start_positions is not None:
            slen, hsz = hidden_states.shape[-2:]
            start_positions = start_positions[:, None, None].expand(-1, -1, hsz)
            start_states = hidden_states.gather(-2, start_positions)
            start_states = start_states.expand(-1, slen, -1)
        x = self.dense_0(torch.cat([hidden_states, start_states], dim=-1))
        x = self.activation(x)
        x = self.LayerNorm(x)
        x = self.dense_1(x).squeeze(-1)
        if p_mask is not None:
            if p_mask.dtype == torch.float16:
                x = x * (1 - p_mask) - 65500 * p_mask
            else:
                x = x * (1 - p_mask) - 1e+30 * p_mask
        return x

Compute SQuAD end logits from sequence hidden states.

Args:
    config ([`XLMConfig`]):
        The config used by the model, will be used to grab the `hidden_size` of the model and the `layer_norm_eps`
        to use.


def tridiagonal_matmul(diagonals, rhs, diagonals_format='compact', name=None):
    if diagonals_format == 'compact':
        superdiag = diagonals[..., 0, :]
        maindiag = diagonals[..., 1, :]
        subdiag = diagonals[..., 2, :]
    elif diagonals_format == 'sequence':
        superdiag, maindiag, subdiag = diagonals
    elif diagonals_format == 'matrix':
        m1 = tensor_shape.dimension_value(diagonals.shape[-1])
        m2 = tensor_shape.dimension_value(diagonals.shape[-2])
        if m1 and m2 and (m1 != m2):
            raise ValueError('Expected last two dimensions of diagonals to be same, got {} and {}'.format(m1, m2))
        diags = array_ops.matrix_diag_part(diagonals, k=(-1, 1), padding_value=0.0, align='LEFT_RIGHT')
        superdiag = diags[..., 0, :]
        maindiag = diags[..., 1, :]
        subdiag = diags[..., 2, :]
    else:
        raise ValueError('Unrecognized diagonals_format: %s' % diagonals_format)
    superdiag = array_ops.expand_dims(superdiag, -2)
    maindiag = array_ops.expand_dims(maindiag, -2)
    subdiag = array_ops.expand_dims(subdiag, -2)
    return linalg_ops.tridiagonal_mat_mul(superdiag, maindiag, subdiag, rhs, name)

Multiplies tridiagonal matrix by matrix.

`diagonals` is representation of 3-diagonal NxN matrix, which depends on
`diagonals_format`.

In `matrix` format, `diagonals` must be a tensor of shape `[..., M, M]`, with
two inner-most dimensions representing the square tridiagonal matrices.
Elements outside of the three diagonals will be ignored.

If `sequence` format, `diagonals` is list or tuple of three tensors:
`[superdiag, maindiag, subdiag]`, each having shape [..., M]. Last element
of `superdiag` first element of `subdiag` are ignored.

In `compact` format the three diagonals are brought together into one tensor
of shape `[..., 3, M]`, with last two dimensions containing superdiagonals,
diagonals, and subdiagonals, in order. Similarly to `sequence` format,
elements `diagonals[..., 0, M-1]` and `diagonals[..., 2, 0]` are ignored.

The `sequence` format is recommended as the one with the best performance.

`rhs` is matrix to the right of multiplication. It has shape `[..., M, N]`.

Example:

```python
superdiag = tf.constant([-1, -1, 0], dtype=tf.float64)
maindiag = tf.constant([2, 2, 2], dtype=tf.float64)
subdiag = tf.constant([0, -1, -1], dtype=tf.float64)
diagonals = [superdiag, maindiag, subdiag]
rhs = tf.constant([[1, 1], [1, 1], [1, 1]], dtype=tf.float64)
x = tf.linalg.tridiagonal_matmul(diagonals, rhs, diagonals_format='sequence')
```

Args:
  diagonals: A `Tensor` or tuple of `Tensor`s describing left-hand sides. The
    shape depends of `diagonals_format`, see description above. Must be
    `float32`, `float64`, `complex64`, or `complex128`.
  rhs: A `Tensor` of shape [..., M, N] and with the same dtype as `diagonals`.
  diagonals_format: one of `sequence`, or `compact`. Default is `compact`.
  name:  A name to give this `Op` (optional).

Returns:
  A `Tensor` of shape [..., M, N] containing the result of multiplication.

Raises:
  ValueError: An unsupported type is provided as input, or when the input
  tensors have incorrect shapes.


def _check_rnn_cell_input_dtypes(inputs):
    for t in nest.flatten(inputs):
        _check_supported_dtypes(t.dtype)

Check whether the input tensors are with supported dtypes.

Default RNN cells only support floats and complex as its dtypes since the
activation function (tanh and sigmoid) only allow those types. This function
will throw a proper error message if the inputs is not in a supported type.

Args:
  inputs: tensor or nested structure of tensors that are feed to RNN cell as
    input or state.

Raises:
  ValueError: if any of the input tensor are not having dtypes of float or
    complex.


def forward(self, pixel_values: Tensor, mask_labels: Optional[List[Tensor]]=None, class_labels: Optional[List[Tensor]]=None, pixel_mask: Optional[Tensor]=None, output_hidden_states: Optional[bool]=None, output_auxiliary_logits: Optional[bool]=None, output_attentions: Optional[bool]=None, return_dict: Optional[bool]=None) -> Mask2FormerForUniversalSegmentationOutput:
    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(pixel_values=pixel_values, pixel_mask=pixel_mask, output_hidden_states=output_hidden_states or self.config.use_auxiliary_loss, output_attentions=output_attentions, return_dict=True)
    loss, loss_dict, auxiliary_logits = (None, None, None)
    class_queries_logits = ()
    for decoder_output in outputs.transformer_decoder_intermediate_states:
        class_prediction = self.class_predictor(decoder_output.transpose(0, 1))
        class_queries_logits += (class_prediction,)
    masks_queries_logits = outputs.masks_queries_logits
    auxiliary_logits = self.get_auxiliary_logits(class_queries_logits, masks_queries_logits)
    if mask_labels is not None and class_labels is not None:
        loss_dict = self.get_loss_dict(masks_queries_logits=masks_queries_logits[-1], class_queries_logits=class_queries_logits[-1], mask_labels=mask_labels, class_labels=class_labels, auxiliary_predictions=auxiliary_logits)
        loss = self.get_loss(loss_dict)
    encoder_hidden_states = None
    pixel_decoder_hidden_states = None
    transformer_decoder_hidden_states = None
    if output_hidden_states:
        encoder_hidden_states = outputs.encoder_hidden_states
        pixel_decoder_hidden_states = outputs.pixel_decoder_hidden_states
        transformer_decoder_hidden_states = outputs.transformer_decoder_hidden_states
    output_auxiliary_logits = self.config.output_auxiliary_logits if output_auxiliary_logits is None else output_auxiliary_logits
    if not output_auxiliary_logits:
        auxiliary_logits = None
    output = Mask2FormerForUniversalSegmentationOutput(loss=loss, class_queries_logits=class_queries_logits[-1], masks_queries_logits=masks_queries_logits[-1], auxiliary_logits=auxiliary_logits, encoder_last_hidden_state=outputs.encoder_last_hidden_state, pixel_decoder_last_hidden_state=outputs.pixel_decoder_last_hidden_state, transformer_decoder_last_hidden_state=outputs.transformer_decoder_last_hidden_state, encoder_hidden_states=encoder_hidden_states, pixel_decoder_hidden_states=pixel_decoder_hidden_states, transformer_decoder_hidden_states=transformer_decoder_hidden_states, attentions=outputs.attentions)
    if not return_dict:
        output = tuple((v for v in output.values() if v is not None))
        if loss is not None:
            output = loss + output
    return output

mask_labels (`List[torch.Tensor]`, *optional*):
    List of mask labels of shape `(num_labels, height, width)` to be fed to a model
class_labels (`List[torch.LongTensor]`, *optional*):
    list of target class labels of shape `(num_labels, height, width)` to be fed to a model. They identify the
    labels of `mask_labels`, e.g. the label of `mask_labels[i][j]` if `class_labels[i][j]`.
output_auxiliary_logits (`bool`, *optional*):
    Whether or not to output auxiliary logits.

Examples:

Instance segmentation example:

```python
>>> from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
>>> from PIL import Image
>>> import requests
>>> import torch

>>> # Load Mask2Former trained on COCO instance segmentation dataset
>>> image_processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-small-coco-instance")
>>> model = Mask2FormerForUniversalSegmentation.from_pretrained(
...     "facebook/mask2former-swin-small-coco-instance"
... )

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

>>> with torch.no_grad():
...     outputs = model(**inputs)

>>> # Model predicts class_queries_logits of shape `(batch_size, num_queries)`
>>> # and masks_queries_logits of shape `(batch_size, num_queries, height, width)`
>>> class_queries_logits = outputs.class_queries_logits
>>> masks_queries_logits = outputs.masks_queries_logits

>>> # Perform post-processing to get instance segmentation map
>>> pred_instance_map = image_processor.post_process_instance_segmentation(
...     outputs, target_sizes=[(image.height, image.width)]
... )[0]
>>> print(pred_instance_map.shape)
torch.Size([480, 640])
```

Semantic segmentation example:
```python
>>> from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
>>> from PIL import Image
>>> import requests
>>> import torch

>>> # Load Mask2Former trained on ADE20k semantic segmentation dataset
>>> image_processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-small-ade-semantic")
>>> model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-small-ade-semantic")

>>> url = (
...     "https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg"
... )
>>> image = Image.open(requests.get(url, stream=True).raw)
>>> inputs = image_processor(image, return_tensors="pt")

>>> with torch.no_grad():
...     outputs = model(**inputs)

>>> # Model predicts class_queries_logits of shape `(batch_size, num_queries)`
>>> # and masks_queries_logits of shape `(batch_size, num_queries, height, width)`
>>> class_queries_logits = outputs.class_queries_logits
>>> masks_queries_logits = outputs.masks_queries_logits

>>> # Perform post-processing to get semantic segmentation map
>>> pred_semantic_map = image_processor.post_process_semantic_segmentation(
...     outputs, target_sizes=[(image.height, image.width)]
... )[0]
>>> print(pred_semantic_map.shape)
torch.Size([512, 683])
```

Panoptic segmentation example:

```python
>>> from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
>>> from PIL import Image
>>> import requests
>>> import torch

>>> # Load Mask2Former trained on CityScapes panoptic segmentation dataset
>>> image_processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-small-cityscapes-panoptic")
>>> model = Mask2FormerForUniversalSegmentation.from_pretrained(
...     "facebook/mask2former-swin-small-cityscapes-panoptic"
... )

>>> url = "https://cdn-media.huggingface.co/Inference-API/Sample-results-on-the-Cityscapes-dataset-The-above-images-show-how-our-method-can-handle.png"
>>> image = Image.open(requests.get(url, stream=True).raw)
>>> inputs = image_processor(image, return_tensors="pt")

>>> with torch.no_grad():
...     outputs = model(**inputs)

>>> # Model predicts class_queries_logits of shape `(batch_size, num_queries)`
>>> # and masks_queries_logits of shape `(batch_size, num_queries, height, width)`
>>> class_queries_logits = outputs.class_queries_logits
>>> masks_queries_logits = outputs.masks_queries_logits

>>> # Perform post-processing to get panoptic segmentation map
>>> pred_panoptic_map = image_processor.post_process_panoptic_segmentation(
...     outputs, target_sizes=[(image.height, image.width)]
... )[0]["segmentation"]
>>> print(pred_panoptic_map.shape)
torch.Size([338, 676])
```


def decode(self, decoder_input_ids, encoder_outputs, encoder_attention_mask: Optional[jnp.ndarray]=None, decoder_attention_mask: Optional[jnp.ndarray]=None, decoder_position_ids: Optional[jnp.ndarray]=None, past_key_values: Optional[dict]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, train: bool=False, params: Optional[dict]=None, dropout_rng: PRNGKey=None):
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.return_dict
    encoder_hidden_states = encoder_outputs[0]
    if encoder_attention_mask is None:
        batch_size, sequence_length = encoder_hidden_states.shape[:2]
        encoder_attention_mask = jnp.ones((batch_size, sequence_length))
    batch_size, sequence_length = decoder_input_ids.shape
    if decoder_attention_mask is None:
        decoder_attention_mask = jnp.ones((batch_size, sequence_length))
    if decoder_position_ids is None:
        if past_key_values is not None:
            raise ValueError('Make sure to provide `decoder_position_ids` when passing `past_key_values`.')
        decoder_position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))
    rngs = {}
    if dropout_rng is not None:
        rngs['dropout'] = dropout_rng
    inputs = {'params': params or self.params}
    if past_key_values:
        inputs['cache'] = past_key_values
        mutable = ['cache']
    else:
        mutable = False

    def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):
        decoder_module = module._get_decoder_module()
        outputs = decoder_module(decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs)
        hidden_states = outputs[0]
        if self.config.tie_word_embeddings:
            shared_embedding = module.model.variables['params']['shared']['embedding']
            lm_logits = module.lm_head.apply({'params': {'kernel': shared_embedding.T}}, hidden_states)
        else:
            lm_logits = module.lm_head(hidden_states)
        lm_logits += module.final_logits_bias
        return (lm_logits, outputs)
    outputs = self.module.apply(inputs, decoder_input_ids=jnp.array(decoder_input_ids, dtype='i4'), decoder_attention_mask=jnp.array(decoder_attention_mask, dtype='i4'), decoder_position_ids=jnp.array(decoder_position_ids, dtype='i4'), encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=jnp.array(encoder_attention_mask, dtype='i4'), output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, deterministic=not train, rngs=rngs, mutable=mutable, method=_decoder_forward)
    if past_key_values is None:
        lm_logits, decoder_outputs = outputs
    else:
        (lm_logits, decoder_outputs), past = outputs
    if return_dict:
        outputs = FlaxCausalLMOutputWithCrossAttentions(logits=lm_logits, hidden_states=decoder_outputs.hidden_states, attentions=decoder_outputs.attentions, cross_attentions=decoder_outputs.cross_attentions)
    else:
        outputs = (lm_logits,) + decoder_outputs[1:]
    if past_key_values is not None and return_dict:
        outputs['past_key_values'] = unfreeze(past['cache'])
        return outputs
    elif past_key_values is not None and (not return_dict):
        outputs = outputs[:1] + (unfreeze(past['cache']),) + outputs[1:]
    return outputs

Returns:

Example:

```python
>>> import jax.numpy as jnp
>>> from transformers import AutoTokenizer, FlaxBlenderbotForConditionalGeneration

>>> model = FlaxBlenderbotForConditionalGeneration.from_pretrained("facebook/blenderbot-400M-distill")
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot-400M-distill")

>>> text = "My friends are cool but they eat too many carbs."
>>> inputs = tokenizer(text, max_length=1024, return_tensors="jax")
>>> encoder_outputs = model.encode(**inputs)

>>> decoder_start_token_id = model.config.decoder_start_token_id
>>> decoder_input_ids = jnp.ones((inputs.input_ids.shape[0], 1), dtype="i4") * decoder_start_token_id

>>> outputs = model.decode(decoder_input_ids, encoder_outputs)
>>> logits = outputs.logits
```


def load_variable(ckpt_dir_or_file, name):
    if name.endswith(':0'):
        name = name[:-2]
    reader = load_checkpoint(ckpt_dir_or_file)
    return reader.get_tensor(name)

Returns the tensor value of the given variable in the checkpoint.

When the variable name is unknown, you can use `tf.train.list_variables` to
inspect all the variable names.

Example usage:

```python
import tensorflow as tf
a = tf.Variable(1.0)
b = tf.Variable(2.0)
ckpt = tf.train.Checkpoint(var_list={'a': a, 'b': b})
ckpt_path = ckpt.save('tmp-ckpt')
var= tf.train.load_variable(
    ckpt_path, 'var_list/a/.ATTRIBUTES/VARIABLE_VALUE')
print(var)  # 1.0
```

Args:
  ckpt_dir_or_file: Directory with checkpoints file or path to checkpoint.
  name: Name of the variable to return.

Returns:
  A numpy `ndarray` with a copy of the value of this variable.


def convert(model_flags: _model_flags_pb2.ModelFlags, conversion_flags: _conversion_flags_pb2.ConverterFlags, input_data_str: Optional[str]=None, debug_info_str: Optional[str]=None):
    try:
        return wrap_converter.wrapped_convert(model_flags.SerializeToString(), conversion_flags.SerializeToString(), input_data_str, debug_info_str)
    except Exception as e:
        converter_error = ConverterError(str(e))
        for error_data in _metrics_wrapper.retrieve_collected_errors():
            converter_error.append_error(error_data)
            if error_data.error_code == converter_error_data_pb2.ConverterErrorData.ERROR_STATEFUL_PARTITIONED_CALL_IN_FINAL_IR and (not conversion_flags.guarantee_all_funcs_one_use):
                conversion_flags.guarantee_all_funcs_one_use = True
                return convert(model_flags, conversion_flags, input_data_str, debug_info_str)
        raise converter_error

Converts `input_data_str` to a TFLite model.

Args:
  model_flags: Proto describing model properties, see `model_flags.proto`.
  conversion_flags: Proto describing conversion properties, see
    `compiler/mlir/lite/converter_flags.proto`.
  input_data_str: Input data in serialized form (e.g. a graphdef is common, or
    it can be hlo text or proto)
  debug_info_str: Serialized `GraphDebugInfo` proto describing logging
    information.

Returns:
  Converted model in serialized form (e.g. a TFLITE model is common).
Raises:
  ConverterError: When conversion fails in TFLiteConverter, usually due to
    ops not being supported.


class Phi4MultimodalConfig(Phi3Config):
    sub_configs = {'audio_config': Phi4MultimodalAudioConfig, 'vision_config': Phi4MultimodalVisionConfig}

    def __init__(self, vocab_size=200064, hidden_size=3072, intermediate_size=8192, num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=8, resid_pdrop=0.0, embd_pdrop=0.0, attention_dropout=0.0, hidden_act='silu', max_position_embeddings=131072, initializer_range=0.02, rms_norm_eps=1e-05, use_cache=True, tie_word_embeddings=False, rope_theta=10000.0, rope_scaling=None, partial_rotary_factor=1, bos_token_id=199999, eos_token_id=[199999, 200020], pad_token_id=199999, original_max_position_embeddings=4096, sliding_window=None, vision_config=None, audio_config=None, **kwargs):
        super().__init__(vocab_size=vocab_size, hidden_size=hidden_size, intermediate_size=intermediate_size, num_hidden_layers=num_hidden_layers, num_attention_heads=num_attention_heads, num_key_value_heads=num_key_value_heads, resid_pdrop=resid_pdrop, embd_pdrop=embd_pdrop, attention_dropout=attention_dropout, hidden_act=hidden_act, max_position_embeddings=max_position_embeddings, initializer_range=initializer_range, rms_norm_eps=rms_norm_eps, use_cache=use_cache, tie_word_embeddings=tie_word_embeddings, rope_theta=rope_theta, rope_scaling=rope_scaling, partial_rotary_factor=partial_rotary_factor, bos_token_id=bos_token_id, eos_token_id=eos_token_id, pad_token_id=pad_token_id, original_max_position_embeddings=original_max_position_embeddings, sliding_window=sliding_window, **kwargs)
        if isinstance(vision_config, dict):
            vision_config = Phi4MultimodalVisionConfig(**vision_config)
        elif vision_config is None:
            Phi4MultimodalVisionConfig()
        self.vision_config = vision_config
        if isinstance(audio_config, dict):
            audio_config = Phi4MultimodalAudioConfig(**audio_config)
        elif vision_config is None:
            audio_config = Phi4MultimodalAudioConfig()
        self.audio_config = audio_config

This is the configuration class to store the configuration of a [`Phi4MultimodalModel`]. It is used to instantiate a
Phi4Multimodal model according to the specified arguments, defining the model architecture. Instantiating a configuration
with the defaults will yield a similar configuration to that of the
[microsoft/Phi-4-multimodal-instruct](https://huggingface.co/microsoft/Phi-4-multimodal-instruct) architecture.

Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.

Args:
    vocab_size (`int`, *optional*, defaults to 200064):
        Vocabulary size of the Phi-3 model. Defines the number of different tokens that can be represented by the
        `inputs_ids` passed when calling [`Phi3Model`].
    hidden_size (`int`, *optional*, defaults to 3072):
        Dimension of the hidden representations.
    intermediate_size (`int`, *optional*, defaults to 8192):
        Dimension of the MLP representations.
    num_hidden_layers (`int`, *optional*, defaults to 32):
        Number of hidden layers in the Transformer decoder.
    num_attention_heads (`int`, *optional*, defaults to 32):
        Number of attention heads for each attention layer in the Transformer decoder.
    num_key_value_heads (`int`, *optional*, defaults to 8):
        This is the number of key_value heads that should be used to implement Grouped Query Attention. If
        `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
        `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
        converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
        by meanpooling all the original heads within that group. For more details, check out [this
        paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to
        `num_attention_heads`.
    resid_pdrop (`float`, *optional*, defaults to 0.0):
        Dropout probability for mlp outputs.
    embd_pdrop (`int`, *optional*, defaults to 0.0):
        The dropout ratio for the embeddings.
    attention_dropout (`float`, *optional*, defaults to 0.0):
        The dropout ratio after computing the attention scores.
    hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
        The non-linear activation function (function or string) in the decoder.
    max_position_embeddings (`int`, *optional*, defaults to 131072):
        The maximum sequence length that this model might ever be used with.
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    rms_norm_eps (`float`, *optional*, defaults to 1e-05):
        The epsilon value used for the RMSNorm.
    use_cache (`bool`, *optional*, defaults to `True`):
        Whether or not the model should return the last key/values attentions (not used by all models). Only
        relevant if `config.is_decoder=True`. Whether to tie weight embeddings or not.
    tie_word_embeddings (`bool`, *optional*, defaults to `False`):
        Whether to tie weight embeddings
    rope_theta (`float`, *optional*, defaults to 10000.0):
        The base period of the RoPE embeddings.
    rope_scaling (`dict`, *optional*):
        The scaling strategy for the RoPE embeddings. If `None`, no scaling is applied. If a dictionary, it must
        contain the following keys: `type`, `short_factor` and `long_factor`. The `type` must be `longrope` and
        the `short_factor` and `long_factor` must be lists of numbers with the same length as the hidden size
        divided by the number of attention heads divided by 2.
    partial_rotary_factor (`float`, *optional*, defaults to `1.0`):
        Percentage of the query and keys which will have rotary embedding. Must be between 0.0 and 1.0.
    bos_token_id (`int`, *optional*, defaults to 199999):
        The id of the "beginning-of-sequence" token.
    eos_token_id (`int` or `list[int]`, *optional*, defaults to `[199999, 200020]`):
        The id of the "end-of-sequence" token.
    pad_token_id (`int`, *optional*, defaults to 199999):
        The id of the padding token.
    original_max_position_embeddings (`int`, *optional*, defaults to 4096):
        The maximum sequence length that this model was trained with. This is used to determine the size of the
        original RoPE embeddings when using long scaling.
    sliding_window (`int`, *optional*):
        Sliding window attention window size. If `None`, no sliding window is applied.
    vision_config (`Phi4MultimodalVisionConfig` or `dict`, *optional*):
        The vision config for the underlying image embedding model. If not provided, will default to the configuration
        used to instantiate a model similar in architecture as
        [microsoft/Phi-4-multimodal-instruct](https://huggingface.co/microsoft/Phi-4-multimodal-instruct).
    audio_config (`Phi4MultimodalAudioConfig` or `dict`, *optional*):
        The audio config for the underlying audio embedding model. If not provided, will default to the configuration
        used to instantiate a model similar in architecture as
        [microsoft/Phi-4-multimodal-instruct](https://huggingface.co/microsoft/Phi-4-multimodal-instruct).

Example:

```python
>>> from transformers import Phi4MultimodalModel, Phi4MultimodalConfig

>>> # Initializing a Phi4Multimodal style configuration
>>> configuration = Phi4MultimodalConfig.from_pretrained("microsoft/Phi-4-multimodal-instruct")

>>> # Initializing a model from the configuration
>>> model = Phi4MultimodalModel(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config
```


class MegaConfig(PretrainedConfig):
    model_type = 'mega'

    def __init__(self, vocab_size=30522, hidden_size=128, num_hidden_layers=4, intermediate_size=256, ema_projection_size=16, bidirectional=True, shared_representation_size=64, use_chunking=False, chunk_size=-1, truncation=None, normalize_before_mega=True, normalization_type='scalenorm', norm_affine=True, activation='silu', attention_activation='softmax', dropout_prob=0.1, hidden_dropout_prob=0.1, attention_probs_dropout_prob=0.1, use_feature_dropout=False, use_normalized_ffn=True, nffn_hidden_size=256, normalize_before_ffn=True, nffn_activation_dropout_prob=0.1, max_positions=2048, add_token_type_embeddings=False, type_vocab_size=2, initializer_range=0.02, ema_delta_alpha_range=0.2, ema_beta_range=0.02, ema_gamma_omega_range=1.0, pad_token_id=1, bos_token_id=0, eos_token_id=2, relative_positional_bias='rotary', classifier_dropout=None, use_cache=True, add_lm_hidden_dense_layer=True, **kwargs):
        super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.activation = activation
        self.attention_activation = attention_activation
        self.intermediate_size = intermediate_size
        self.ema_projection_size = ema_projection_size
        self.bidirectional = bidirectional
        self.shared_representation_size = shared_representation_size
        self.use_chunking = use_chunking
        self.chunk_size = chunk_size
        self.truncation = truncation
        self.normalize_before_mega = normalize_before_mega
        self.normalization_type = normalization_type
        self.norm_affine = norm_affine
        self.dropout_prob = dropout_prob
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.use_feature_dropout = use_feature_dropout
        self.use_normalized_ffn = use_normalized_ffn
        self.nffn_hidden_size = nffn_hidden_size
        self.normalize_before_ffn = normalize_before_ffn
        self.nffn_activation_dropout_prob = nffn_activation_dropout_prob
        self.max_positions = max_positions
        self.add_token_type_embeddings = add_token_type_embeddings
        self.type_vocab_size = type_vocab_size
        self.initializer_range = initializer_range
        self.ema_delta_alpha_range = ema_delta_alpha_range
        self.ema_beta_range = ema_beta_range
        self.ema_gamma_omega_range = ema_gamma_omega_range
        self.relative_positional_bias = relative_positional_bias
        self.use_cache = use_cache
        self.classifier_dropout = classifier_dropout
        self.add_lm_hidden_dense_layer = add_lm_hidden_dense_layer
        self.num_attention_heads = 1

This is the configuration class to store the configuration of a [`MegaModel`]. It is used to instantiate a Mega
model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
defaults will yield a similar configuration to that of the Mega
[mnaylor/mega-base-wikitext](https://huggingface.co/mnaylor/mega-base-wikitext) architecture.

Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.

Args:
    vocab_size (`int`, *optional*, defaults to 30522):
        Vocabulary size of the Mega model. Defines the number of different tokens that can be represented by the
        `inputs_ids` passed when calling [`MegaModel`].
    hidden_size (`int`, *optional*, defaults to 128):
        Dimensionality of the encoder layers and the pooler layer.
    num_hidden_layers (`int`, *optional*, defaults to 4):
        Number of hidden layers in the Mega encoder.
    intermediate_size (`int`, *optional*, defaults to 256):
        Dimensionality of the hidden size (self-attention value projection) within the Mega encoder
    ema_projection_size (`int`, *optional*, defaults to 16):
        Dimensionality of the MegaMultiDimensionDampedEma
    bidirectional (`bool`, *optional*, defaults to `True`):
        Whether the MegaMultiDimensionDampedEma used in Mega's self-attention should work bidirectionally (`True`)
        or unidirectionally (`False`). Bidirectional EMA is incompatible with causal decoding, so this should be
        False if you intend to use the model as a decoder.
    shared_representation_size (`int`, *optional*, defaults to 64):
        Dimensionality of the linear projection for shared representation of self-attention queries and keys
    use_chunking (`bool`, *optional*, defaults to `False`):
        Whether to chunk inputs for linear self-attention complexity (described as Mega-chunk in the paper)
    chunk_size (`int`, *optional*, defaults to -1):
        If `use_chunking` is set to `True`, determines the size of the chunks to apply to the input sequence. If
        chunking is used, input sequences must be padded to a multiple of `chunk_size`
    truncation (`int`, *optional*):
        If specified, the sequence length for which to truncate MegaMultiDimensionDampedEma
    normalize_before_mega (`bool`, *optional*, defaults to `True`):
        Whether to normalize before (`True`) or after (`False`) passing through Mega encoder blocks
    normalization_type (`str`, *optional*, defaults to `"scalenorm"`):
        Type of normalization to use in Mega encoder blocks. Choose one of `"scalenorm"`, `"layernorm"`,
        `"rmsnorm"`, `"batchnorm"`, or `"syncbatchnorm"` (GPU required for syncbatchnorm)
    norm_affine (`bool`, *optional*, defaults to `True`):
        If `True`, applies a parameterized affine transformation to inputs during normalization
    activation (`str`, *optional*, defaults to `"silu"`):
        Activation function to apply within Mega encoder blocks. Choose one of `"silu"`, `"relu"`, `"linear"`,
        `"gelu"`, or `"gelu_accurate"`
    attention_activation (`str`, *optional*, defaults to `"softmax"`):
        Activation function to apply for single-headed self-attention (a la Transformer). Choose one of
        `"softmax"`, `"laplace"`, or `"relu2"`
    dropout_prob (`float`, *optional*, defaults to 0.1):
        The dropout probability for EMA self-attention
    hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
        The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
    attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
        The dropout ratio for the attention probabilities.
    use_feature_dropout (`bool`, *optional*, defaults to `False`):
        Whether to use feature-based (`True`) or standard dropout (`False`)
    use_normalized_ffn (`bool`, *optional*, defaults to `True`):
        Whether to use the normalized feed-forward sub-layer in Mega blocks (`True`) or pass Mega encoder output
        as-is (`False`)
    nffn_hidden_size (`int`, *optional*, defaults to 256):
        If using the normalized feed-forward network (NFFN) layer within Mega (`use_normalized_ffn = True`), this
        is the hidden size of the NFFN
    normalize_before_ffn (`bool`, *optional*, defaults to `True`):
        Whether to normalize before (`True`) or after (`False`) the feed-forward portion of NFFN
    nffn_activation_dropout_prob (`float`, *optional*, defaults to 0.1):
        The dropout ratio for the NFFN component.
    max_positions (`int`, *optional*, defaults to 2048):
        The maximum sequence length to use for positional representations. For `"simple"` relative positional bias,
        this is a hard limit on input length; `"rotary"` relative positional bias will extrapolate to longer
        sequences
    add_token_type_embeddings (`bool`, *optional*, defaults to `True`):
        Whether to account for token types in embeddings. Left as optional to maintain compatibility with original
        implementation while adding support for token types.
    type_vocab_size (`int`, *optional*, defaults to 2):
        The vocabulary size of the `token_type_ids` passed when calling [`MegaModel`]. Only used if
        `add_token_type_embeddings = True`
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    ema_delta_alpha_range (`float`, *optional*, defaults to 0.2):
        The standard deviation for initializing the delta (damping factor) and alpha (decay factor) parameters in
        MegaMultiDimensionDampedEma.
    ema_beta_range (`float`, *optional*, defaults to 0.02):
        The standard deviation for initializing the beta parameter (expansion matrix) in
        MegaMultiDimensionDampedEma.
    ema_gamma_omega_range (`float`, *optional*, defaults to 1.0):
        The standard deviation for initializing the gamma (projection matrix) and omega (residual weight)
        parameters in MultiDimensionEMA.
    relative_positional_bias (`str`, *optional*, defaults to `"rotary"`):
        Type of relative positional encoding. Choose one of `"rotary"` or `"simple"`. If `"simple"` is selected,
        `max_positions` is used as a limit on input size, while `"rotary"` extrapolates beyond `max_positions`.
    is_decoder (`bool`, *optional*, defaults to `False`):
        Whether the model is used as a decoder or not. If `False`, the model is used as an encoder.
    use_cache (`bool`, *optional*, defaults to `True`):
        Whether or not the model should return the last key/values attentions (not used by all models). Only
        relevant if `config.is_decoder=True`.
    classifier_dropout (`float`, *optional*):
        The dropout ratio for the classification head.
    add_lm_hidden_dense_layer (`bool`, *optional*, defaults to `True`):
        Whether to include a hidden layer for projection between encoder outputs and LM heads (`True`) or pass
        hidden states directly to LM head (`False`). Remains optional for compatibility with original
        implementation

Examples:

```python
>>> from transformers import MegaConfig, MegaModel

>>> # Initializing a Mega configuration
>>> configuration = MegaConfig()

>>> # Initializing a model (with random weights) from the configuration
>>> model = MegaModel(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config
```


def conjugate_gradient(operator, rhs, preconditioner=None, x=None, tol=1e-05, max_iter=20, name='conjugate_gradient'):
    if not (operator.is_self_adjoint and operator.is_positive_definite):
        raise ValueError('Expected a self-adjoint, positive definite operator.')
    cg_state = collections.namedtuple('CGState', ['i', 'x', 'r', 'p', 'gamma'])

    def stopping_criterion(i, state):
        return math_ops.logical_and(i < max_iter, math_ops.reduce_any(linalg.norm(state.r, axis=-1) > tol))

    def dot(x, y):
        return array_ops.squeeze(math_ops.matvec(x[..., array_ops.newaxis], y, adjoint_a=True), axis=-1)

    def cg_step(i, state):
        z = math_ops.matvec(operator, state.p)
        alpha = state.gamma / dot(state.p, z)
        x = state.x + alpha[..., array_ops.newaxis] * state.p
        r = state.r - alpha[..., array_ops.newaxis] * z
        if preconditioner is None:
            q = r
        else:
            q = preconditioner.matvec(r)
        gamma = dot(r, q)
        beta = gamma / state.gamma
        p = q + beta[..., array_ops.newaxis] * state.p
        return (i + 1, cg_state(i + 1, x, r, p, gamma))
    with ops.name_scope(name):
        broadcast_shape = array_ops.broadcast_dynamic_shape(array_ops.shape(rhs)[:-1], operator.batch_shape_tensor())
        if preconditioner is not None:
            broadcast_shape = array_ops.broadcast_dynamic_shape(broadcast_shape, preconditioner.batch_shape_tensor())
        broadcast_rhs_shape = array_ops.concat([broadcast_shape, [array_ops.shape(rhs)[-1]]], axis=-1)
        r0 = array_ops.broadcast_to(rhs, broadcast_rhs_shape)
        tol *= linalg.norm(r0, axis=-1)
        if x is None:
            x = array_ops.zeros(broadcast_rhs_shape, dtype=rhs.dtype.base_dtype)
        else:
            r0 = rhs - math_ops.matvec(operator, x)
        if preconditioner is None:
            p0 = r0
        else:
            p0 = math_ops.matvec(preconditioner, r0)
        gamma0 = dot(r0, p0)
        i = constant_op.constant(0, dtype=dtypes.int32)
        state = cg_state(i=i, x=x, r=r0, p=p0, gamma=gamma0)
        _, state = while_loop.while_loop(stopping_criterion, cg_step, [i, state])
        return cg_state(state.i, x=state.x, r=state.r, p=state.p, gamma=state.gamma)

Conjugate gradient solver.

Solves a linear system of equations `A*x = rhs` for self-adjoint, positive
definite matrix `A` and right-hand side vector `rhs`, using an iterative,
matrix-free algorithm where the action of the matrix A is represented by
`operator`. The iteration terminates when either the number of iterations
exceeds `max_iter` or when the residual norm has been reduced to `tol`
times its initial value, i.e. \\(||rhs - A x_k|| <= tol ||rhs||\\).

Args:
  operator: A `LinearOperator` that is self-adjoint and positive definite.
  rhs: A possibly batched vector of shape `[..., N]` containing the right-hand
    size vector.
  preconditioner: A `LinearOperator` that approximates the inverse of `A`.
    An efficient preconditioner could dramatically improve the rate of
    convergence. If `preconditioner` represents matrix `M`(`M` approximates
    `A^{-1}`), the algorithm uses `preconditioner.apply(x)` to estimate
    `A^{-1}x`. For this to be useful, the cost of applying `M` should be
    much lower than computing `A^{-1}` directly.
  x: A possibly batched vector of shape `[..., N]` containing the initial
    guess for the solution.
  tol: A float scalar convergence tolerance.
  max_iter: An integer giving the maximum number of iterations.
  name: A name scope for the operation.

Returns:
  output: A namedtuple representing the final state with fields:
    - i: A scalar `int32` `Tensor`. Number of iterations executed.
    - x: A rank-1 `Tensor` of shape `[..., N]` containing the computed
        solution.
    - r: A rank-1 `Tensor` of shape `[.., M]` containing the residual vector.
    - p: A rank-1 `Tensor` of shape `[..., N]`. `A`-conjugate basis vector.
    - gamma: \\(r \dot M \dot r\\), equivalent to  \\(||r||_2^2\\) when
      `preconditioner=None`.


class TFCLIPEncoder(keras.layers.Layer):

    def __init__(self, config: CLIPConfig, **kwargs):
        super().__init__(**kwargs)
        self.layers = [TFCLIPEncoderLayer(config, name=f'layers_._{i}') for i in range(config.num_hidden_layers)]

    def call(self, hidden_states: tf.Tensor, attention_mask: tf.Tensor, causal_attention_mask: tf.Tensor, output_attentions: bool, output_hidden_states: bool, return_dict: bool, training: bool=False) -> Union[TFBaseModelOutput, Tuple[tf.Tensor]]:
        all_hidden_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None
        for i, layer_module in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)
            layer_outputs = layer_module(hidden_states=hidden_states, attention_mask=attention_mask, causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, training=training)
            hidden_states = layer_outputs[0]
            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)
        if not return_dict:
            return tuple((v for v in [hidden_states, all_hidden_states, all_attentions] if v is not None))
        return TFBaseModelOutput(last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions)

    def build(self, input_shape=None):
        if self.built:
            return
        self.built = True
        if getattr(self, 'layers', None) is not None:
            for layer in self.layers:
                with tf.name_scope(layer.name):
                    layer.build(None)

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

Args:
    config: CLIPConfig


def init_op(self):
    return self._init_op

Return the Init Op used by the supervisor.

Returns:
  An Op or `None`.


def divide_no_nan(x1, x2):
    if any_symbolic_tensors((x1, x2)):
        return DivideNoNan().symbolic_call(x1, x2)
    return backend.numpy.divide_no_nan(x1, x2)

Safe element-wise division which returns 0 where the denominator is 0.

Args:
    x1: First input tensor.
    x2: Second input tensor.

Returns:
    The quotient `x1/x2`, element-wise, with zero where x2 is zero.


def __call__(self, *args: Union[str, 'Image.Image', List['Image.Image'], List[str]], **kwargs: Any) -> List[Any]:
    return super().__call__(*args, **kwargs)

Extract the features of the input(s).

Args:
    images (`str`, `List[str]`, `PIL.Image` or `List[PIL.Image]`):
        The pipeline handles three types of images:

        - A string containing a http link pointing to an image
        - A string containing a local path to an image
        - An image loaded in PIL directly

        The pipeline accepts either a single image or a batch of images, which must then be passed as a string.
        Images in a batch must all be in the same format: all as http links, all as local paths, or all as PIL
        images.
    timeout (`float`, *optional*, defaults to None):
        The maximum time in seconds to wait for fetching images from the web. If None, no timeout is used and
        the call may block forever.
Return:
    A nested list of `float`: The features computed by the model.


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

Args:
    hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
    original_hidden_states (`torch.FloatTensor`): word embedding output that will be concatenated with
    hidden activations to form the input of the shared transformer layer.
    layer_idx (`int`): layer number.
    attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
        `(batch, sequence_length)` where padding elements are indicated by 0.
    past_key_value (`ZambaHybridDynamicCache`, *optional*): cached past key and value projection states
    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.
    use_cache (`bool`, *optional*):
        If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
        (see `past_key_values`).
    cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
        Indices depicting the position of the input sequence tokens in the sequence.


def from_json(cls, json_value: Any, *, value_spec: Optional[pg_typing.List]=None, allow_partial: bool=False, root_path: Optional[utils.KeyPath]=None, **kwargs) -> 'List':
    return cls([base.from_json(v, root_path=utils.KeyPath(i, root_path), allow_partial=allow_partial, **kwargs) for i, v in enumerate(json_value)], value_spec=value_spec, root_path=root_path, allow_partial=allow_partial)

Class method that load an symbolic List from a JSON value.

Example::

    l = List.from_json([{
        '_type': '__main__.Foo',
        'f1': 1,
        'f2': {
          'f21': True
        }
      },
      1
    ])

    assert l.value_spec is None
    # Okay:
    l.append('abc')

    # [0].f2 is bound by class Foo's field 'f2' definition
    # (assuming it defines a schema for the Dict field).
    assert l[0].f2.value_spec is not None

    # Not okay:
    l[0].f2.abc = 1

Args:
  json_value: Input JSON value, only JSON list is acceptable.
  value_spec: An optional `pg.typing.List` object as the schema for the
    list.
  allow_partial: Whether to allow elements of the list to be partial.
  root_path: KeyPath of loaded object in its object tree.
  **kwargs: Allow passing through keyword arguments that are not applicable.

Returns:
  A schema-less symbolic list, but its items maybe symbolic.


def MonitoredTrainingSession(master='', is_chief=True, checkpoint_dir=None, scaffold=None, hooks=None, chief_only_hooks=None, save_checkpoint_secs=USE_DEFAULT, save_summaries_steps=USE_DEFAULT, save_summaries_secs=USE_DEFAULT, config=None, stop_grace_period_secs=120, log_step_count_steps=100, max_wait_secs=7200, save_checkpoint_steps=USE_DEFAULT, summary_dir=None, save_graph_def=True):
    if save_summaries_steps == USE_DEFAULT and save_summaries_secs == USE_DEFAULT:
        save_summaries_steps = 100
        save_summaries_secs = None
    elif save_summaries_secs == USE_DEFAULT:
        save_summaries_secs = None
    elif save_summaries_steps == USE_DEFAULT:
        save_summaries_steps = None
    if save_checkpoint_steps == USE_DEFAULT and save_checkpoint_secs == USE_DEFAULT:
        save_checkpoint_steps = None
        save_checkpoint_secs = 600
    elif save_checkpoint_secs == USE_DEFAULT:
        save_checkpoint_secs = None
    elif save_checkpoint_steps == USE_DEFAULT:
        save_checkpoint_steps = None
    scaffold = scaffold or Scaffold()
    worker_context = distribute_coordinator_context.get_current_worker_context()
    if worker_context:
        return _create_monitored_session_with_worker_context(worker_context, scaffold, checkpoint_dir=checkpoint_dir, hooks=hooks, chief_only_hooks=chief_only_hooks, save_checkpoint_secs=save_checkpoint_secs, save_summaries_steps=save_summaries_steps, save_summaries_secs=save_summaries_secs, config=config, stop_grace_period_secs=stop_grace_period_secs, log_step_count_steps=log_step_count_steps, max_wait_secs=max_wait_secs, save_checkpoint_steps=save_checkpoint_steps, summary_dir=summary_dir, save_graph_def=save_graph_def)
    if not is_chief:
        session_creator = WorkerSessionCreator(scaffold=scaffold, master=master, config=config, max_wait_secs=max_wait_secs)
        return MonitoredSession(session_creator=session_creator, hooks=hooks or [], stop_grace_period_secs=stop_grace_period_secs)
    all_hooks = []
    if chief_only_hooks:
        all_hooks.extend(chief_only_hooks)
    session_creator = ChiefSessionCreator(scaffold=scaffold, checkpoint_dir=checkpoint_dir, master=master, config=config)
    summary_dir = summary_dir or checkpoint_dir
    if summary_dir:
        if log_step_count_steps and log_step_count_steps > 0:
            all_hooks.append(basic_session_run_hooks.StepCounterHook(output_dir=summary_dir, every_n_steps=log_step_count_steps))
        if save_summaries_steps and save_summaries_steps > 0 or (save_summaries_secs and save_summaries_secs > 0):
            all_hooks.append(basic_session_run_hooks.SummarySaverHook(scaffold=scaffold, save_steps=save_summaries_steps, save_secs=save_summaries_secs, output_dir=summary_dir))
    if checkpoint_dir:
        if save_checkpoint_secs and save_checkpoint_secs > 0 or (save_checkpoint_steps and save_checkpoint_steps > 0):
            all_hooks.append(basic_session_run_hooks.CheckpointSaverHook(checkpoint_dir, save_steps=save_checkpoint_steps, save_secs=save_checkpoint_secs, scaffold=scaffold, save_graph_def=save_graph_def))
    if hooks:
        all_hooks.extend(hooks)
    return MonitoredSession(session_creator=session_creator, hooks=all_hooks, stop_grace_period_secs=stop_grace_period_secs)

Creates a `MonitoredSession` for training.

For a chief, this utility sets proper session initializer/restorer. It also
creates hooks related to checkpoint and summary saving. For workers, this
utility sets proper session creator which waits for the chief to
initialize/restore. Please check `tf.compat.v1.train.MonitoredSession` for
more
information.

@compatibility(TF2)
This API is not compatible with eager execution and `tf.function`. To migrate
to TF2, rewrite the code to be compatible with eager execution. Check the
[migration
guide](https://www.tensorflow.org/guide/migrate#1_replace_v1sessionrun_calls)
on replacing `Session.run` calls. In Keras, session hooks can be replaced by
Callbacks e.g. [logging hook notebook](
https://github.com/tensorflow/docs/blob/master/site/en/guide/migrate/logging_stop_hook.ipynb)
For more details please read [Better
performance with tf.function](https://www.tensorflow.org/guide/function).
@end_compatibility

Args:
  master: `String` the TensorFlow master to use.
  is_chief: If `True`, it will take care of initialization and recovery the
    underlying TensorFlow session. If `False`, it will wait on a chief to
    initialize or recover the TensorFlow session.
  checkpoint_dir: A string.  Optional path to a directory where to restore
    variables.
  scaffold: A `Scaffold` used for gathering or building supportive ops. If not
    specified, a default one is created. It's used to finalize the graph.
  hooks: Optional list of `SessionRunHook` objects.
  chief_only_hooks: list of `SessionRunHook` objects. Activate these hooks if
    `is_chief==True`, ignore otherwise.
  save_checkpoint_secs: The frequency, in seconds, that a checkpoint is saved
    using a default checkpoint saver. If both `save_checkpoint_steps` and
    `save_checkpoint_secs` are set to `None`, then the default checkpoint
    saver isn't used. If both are provided, then only `save_checkpoint_secs`
    is used. Default 600.
  save_summaries_steps: The frequency, in number of global steps, that the
    summaries are written to disk using a default summary saver. If both
    `save_summaries_steps` and `save_summaries_secs` are set to `None`, then
    the default summary saver isn't used. Default 100.
  save_summaries_secs: The frequency, in secs, that the summaries are written
    to disk using a default summary saver.  If both `save_summaries_steps` and
    `save_summaries_secs` are set to `None`, then the default summary saver
    isn't used. Default not enabled.
  config: an instance of `tf.compat.v1.ConfigProto` proto used to configure
    the session. It's the `config` argument of constructor of
    `tf.compat.v1.Session`.
  stop_grace_period_secs: Number of seconds given to threads to stop after
    `close()` has been called.
  log_step_count_steps: The frequency, in number of global steps, that the
    global step/sec is logged.
  max_wait_secs: Maximum time workers should wait for the session to become
    available. This should be kept relatively short to help detect incorrect
    code, but sometimes may need to be increased if the chief takes a while to
    start up.
  save_checkpoint_steps: The frequency, in number of global steps, that a
    checkpoint is saved using a default checkpoint saver. If both
    `save_checkpoint_steps` and `save_checkpoint_secs` are set to `None`, then
    the default checkpoint saver isn't used. If both are provided, then only
    `save_checkpoint_secs` is used. Default not enabled.
  summary_dir: A string.  Optional path to a directory where to save
    summaries. If None, checkpoint_dir is used instead.
  save_graph_def: Whether to save the GraphDef and MetaGraphDef to
    `checkpoint_dir`. The GraphDef is saved after the session is created as
    `graph.pbtxt`. MetaGraphDefs are saved out for every checkpoint as
    `model.ckpt-*.meta`.

Returns:
  A `MonitoredSession` object.


def _on_fail(self, record):
    self.on_fail(record)

Proxy function to guarantee the base implementation of on_fail is
called.

Args:
  record: records.TestResultRecord, a copy of the test record for
      this test, containing all information of the test execution
      including exception objects.


def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    q_type, k_type = (q.dtype, k.dtype)
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    q_embed = q * cos + rotate_half(q) * sin
    k_embed = k * cos + rotate_half(k) * sin
    return (q_embed.to(q_type), k_embed.to(k_type))

Applies Rotary Position Embedding to the query and key tensors.

Args:
    q (`torch.Tensor`): The query tensor.
    k (`torch.Tensor`): The key tensor.
    cos (`torch.Tensor`): The cosine part of the rotary embedding.
    sin (`torch.Tensor`): The sine part of the rotary embedding.
    position_ids (`torch.Tensor`, *optional*):
        Deprecated and unused.
    unsqueeze_dim (`int`, *optional*, defaults to 1):
        The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
        sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
        that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
        k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
        cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
        the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
Returns:
    `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.


def _resize_images(self, x, height_factor, width_factor, data_format, interpolation='nearest'):
    if data_format not in {'channels_last', 'channels_first'}:
        raise ValueError(f'Invalid `data_format` argument: {data_format}')
    if data_format == 'channels_first':
        x = ops.transpose(x, [0, 2, 3, 1])
    if interpolation == 'nearest':
        x = ops.repeat(x, height_factor, axis=1)
        x = ops.repeat(x, width_factor, axis=2)
    else:
        shape = ops.shape(x)
        new_shape = (shape[1] * height_factor, shape[2] * width_factor)
        x = ops.image.resize(x, new_shape, interpolation=interpolation)
    if data_format == 'channels_first':
        x = ops.transpose(x, [0, 3, 1, 2])
    return x

Resizes the images contained in a 4D tensor.

Args:
    x: Tensor or variable to resize.
    height_factor: Positive integer.
    width_factor: Positive integer.
    data_format: One of `"channels_first"`, `"channels_last"`.
    interpolation: A string, one of `"bicubic"`, `"bilinear"`,
    `"lanczos3"`, `"lanczos5"`, or `"nearest"`.

Returns:
    A tensor.


class XLNetSequenceSummary(nn.Module):

    def __init__(self, config: XLNetConfig):
        super().__init__()
        self.summary_type = getattr(config, 'summary_type', 'last')
        if self.summary_type == 'attn':
            raise NotImplementedError
        self.summary = nn.Identity()
        if hasattr(config, 'summary_use_proj') and config.summary_use_proj:
            if hasattr(config, 'summary_proj_to_labels') and config.summary_proj_to_labels and (config.num_labels > 0):
                num_classes = config.num_labels
            else:
                num_classes = config.hidden_size
            self.summary = nn.Linear(config.hidden_size, num_classes)
        activation_string = getattr(config, 'summary_activation', None)
        self.activation: Callable = get_activation(activation_string) if activation_string else nn.Identity()
        self.first_dropout = nn.Identity()
        if hasattr(config, 'summary_first_dropout') and config.summary_first_dropout > 0:
            self.first_dropout = nn.Dropout(config.summary_first_dropout)
        self.last_dropout = nn.Identity()
        if hasattr(config, 'summary_last_dropout') and config.summary_last_dropout > 0:
            self.last_dropout = nn.Dropout(config.summary_last_dropout)

    def forward(self, hidden_states: torch.FloatTensor, cls_index: Optional[torch.LongTensor]=None) -> torch.FloatTensor:
        """
        Compute a single vector summary of a sequence hidden states.

        Args:
            hidden_states (`torch.FloatTensor` of shape `[batch_size, seq_len, hidden_size]`):
                The hidden states of the last layer.
            cls_index (`torch.LongTensor` of shape `[batch_size]` or `[batch_size, ...]` where ... are optional leading dimensions of `hidden_states`, *optional*):
                Used if `summary_type == "cls_index"` and takes the last token of the sequence as classification token.

        Returns:
            `torch.FloatTensor`: The summary of the sequence hidden states.
        """
        if self.summary_type == 'last':
            output = hidden_states[:, -1]
        elif self.summary_type == 'first':
            output = hidden_states[:, 0]
        elif self.summary_type == 'mean':
            output = hidden_states.mean(dim=1)
        elif self.summary_type == 'cls_index':
            if cls_index is None:
                cls_index = torch.full_like(hidden_states[..., :1, :], hidden_states.shape[-2] - 1, dtype=torch.long)
            else:
                cls_index = cls_index.unsqueeze(-1).unsqueeze(-1)
                cls_index = cls_index.expand((-1,) * (cls_index.dim() - 1) + (hidden_states.size(-1),))
            output = hidden_states.gather(-2, cls_index).squeeze(-2)
        elif self.summary_type == 'attn':
            raise NotImplementedError
        output = self.first_dropout(output)
        output = self.summary(output)
        output = self.activation(output)
        output = self.last_dropout(output)
        return output

Compute a single vector summary of a sequence hidden states.

Args:
    config ([`XLNetConfig`]):
        The config used by the model. Relevant arguments in the config class of the model are (refer to the actual
        config class of your model for the default values it uses):

        - **summary_type** (`str`) -- The method to use to make this summary. Accepted values are:

            - `"last"` -- Take the last token hidden state (like XLNet)
            - `"first"` -- Take the first token hidden state (like Bert)
            - `"mean"` -- Take the mean of all tokens hidden states
            - `"cls_index"` -- Supply a Tensor of classification token position (GPT/GPT-2)
            - `"attn"` -- Not implemented now, use multi-head attention

        - **summary_use_proj** (`bool`) -- Add a projection after the vector extraction.
        - **summary_proj_to_labels** (`bool`) -- If `True`, the projection outputs to `config.num_labels` classes
          (otherwise to `config.hidden_size`).
        - **summary_activation** (`Optional[str]`) -- Set to `"tanh"` to add a tanh activation to the output,
          another string or `None` will add no activation.
        - **summary_first_dropout** (`float`) -- Optional dropout probability before the projection and activation.
        - **summary_last_dropout** (`float`)-- Optional dropout probability after the projection and activation.


def binomial(self, shape, counts, probs, dtype=dtypes.int32, name=None):
    dtype = dtypes.as_dtype(dtype)
    with ops.name_scope(name, 'binomial', [shape, counts, probs]) as name:
        counts = ops.convert_to_tensor(counts, name='counts')
        probs = ops.convert_to_tensor(probs, name='probs')
        shape_tensor = _shape_tensor(shape)
        return gen_stateful_random_ops.stateful_random_binomial(self.state.handle, self.algorithm, shape=shape_tensor, counts=counts, probs=probs, dtype=dtype, name=name)

Outputs random values from a binomial distribution.

The generated values follow a binomial distribution with specified count and
probability of success parameters.

Example:

```python
counts = [10., 20.]
# Probability of success.
probs = [0.8]

rng = tf.random.Generator.from_seed(seed=234)
binomial_samples = rng.binomial(shape=[2], counts=counts, probs=probs)

counts = ... # Shape [3, 1, 2]
probs = ...  # Shape [1, 4, 2]
shape = [3, 4, 3, 4, 2]
rng = tf.random.Generator.from_seed(seed=1717)
# Sample shape will be [3, 4, 3, 4, 2]
binomial_samples = rng.binomial(shape=shape, counts=counts, probs=probs)
```

Args:
  shape: A 1-D integer Tensor or Python array. The shape of the output
    tensor.
  counts: Tensor. The counts of the binomial distribution. Must be
    broadcastable with `probs`, and broadcastable with the rightmost
    dimensions of `shape`.
  probs: Tensor. The probability of success for the
    binomial distribution. Must be broadcastable with `counts` and
    broadcastable with the rightmost dimensions of `shape`.
  dtype: The type of the output. Default: tf.int32
  name: A name for the operation (optional).

Returns:
  samples: A Tensor of the specified shape filled with random binomial
    values.  For each i, each samples[i, ...] is an independent draw from
    the binomial distribution on counts[i] trials with probability of
    success probs[i].


def VisitNamedType(self, t):
    if t.name in self._module_map:
        if self._alias_name and '.' in self._alias_name:
            return pytd.Module(name=self._alias_name, module_name=t.name)
        else:
            return t
    module_name, dot, name = t.name.rpartition('.')
    if not dot or self._IsLocalName(module_name):
        return t
    if module_name in self._module_alias_map:
        module_name = self._module_alias_map[module_name]
    try:
        module, cls_prefix = self._LookupModuleRecursive(module_name)
    except KeyError:
        if self._unit and f'{self.name}.{module_name}' in self._unit:
            return t
        raise
    module_name = module.name
    if module_name == self.name:
        return t
    if cls_prefix:
        try:
            maybe_alias = pytd.LookupItemRecursive(module, cls_prefix[:-1])
        except KeyError:
            pass
        else:
            if isinstance(maybe_alias, pytd.Alias) and isinstance(maybe_alias.type, pytd.Module):
                if maybe_alias.type.module_name not in self._module_map:
                    raise KeyError(f'{t.name} refers to unknown module {maybe_alias.name}')
                module = self._module_map[maybe_alias.type.module_name]
                cls_prefix = ''
    name = cls_prefix + name
    try:
        if name == '*':
            self._star_imports.add(module_name)
            item = t
        else:
            item = pytd.LookupItemRecursive(module, name)
    except KeyError as e:
        item = self._ResolveUsingGetattr(module_name, module)
        if item is None:
            item = self._ResolveUsingStarImport(module, name)
            if item is None:
                raise KeyError(f'No {name} in module {module_name}') from e
    if isinstance(item, pytd.Alias):
        lookup_local = LookupLocalTypes()
        lookup_local.unit = module
        new_item = item.Visit(lookup_local)
        if lookup_local.local_names:
            item = new_item
    if not self._in_generic_type and isinstance(item, pytd.Alias):
        item = MaybeSubstituteParameters(item.type) or item
    if isinstance(item, pytd.Constant) and item.name == 'typing_extensions.TypedDict':
        return self.to_type(pytd.NamedType('typing.TypedDict'))
    try:
        return self.to_type(item)
    except NotImplementedError as e:
        raise SymbolLookupError(f'{item} is not a type') from e

Try to look up a NamedType.

Args:
  t: An instance of pytd.NamedType

Returns:
  The same node t.
Raises:
  KeyError: If we can't find a module, or an identifier in a module, or
    if an identifier in a module isn't a class.


def compute_token_logits(sequence_output, temperature, output_weights, output_bias):
    logits = (torch.einsum('bsj,j->bs', sequence_output, output_weights) + output_bias) / temperature
    return logits

Computes logits per token

Args:
    sequence_output (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
        Also known as last_hidden_state. Sequence of hidden-states at the output of the last layer of the model.
    temperature (`float`):
        Temperature for the Bernoulli distribution.
    output_weights (`torch.FloatTensor` of shape `(hidden_size,)`):
        Weights of the linear layer for cell selection.
    output_bias (`torch.FloatTensor` of shape `()`):
        Bias of the linear layer for cell selection

Returns:
    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): Logits per token.


def __init__(self, config, in_features, condition_dim, n_classes=256, bottleneck_factor=2):
    super().__init__()
    bottleneck = (in_features + condition_dim) // bottleneck_factor
    self.mlp = nn.Sequential(nn.Conv2d(in_features + condition_dim, bottleneck, kernel_size=1, stride=1, padding=0), nn.GELU(), nn.Conv2d(bottleneck, 2 + 2, kernel_size=1, stride=1, padding=0), nn.Softplus())
    self.p_eps = 0.0001
    self.max_temp = config.max_temp
    self.min_temp = config.min_temp
    self.log_binomial_transform = LogBinomialSoftmax(n_classes, act=torch.softmax)

Per-pixel MLP followed by a Conditional Log Binomial softmax.

Args:
    in_features (`int`):
        Number of input channels in the main feature.
    condition_dim (`int`):
        Number of input channels in the condition feature.
    n_classes (`int`, *optional*, defaults to 256):
        Number of classes.
    bottleneck_factor (`int`, *optional*, defaults to 2):
        Hidden dim factor.