yobx.xshape.shape_type_compute#

Computes the output shape for broadcasting operators (e.g. Add, Mul, Where).

The function follows NumPy/ONNX broadcasting rules. Shapes are right-aligned and each pair of dimensions (a, b) is resolved according to the table below:

`a`	`b`	Result	Side effect
int (any)	int (any)	max	none
int `n ≠ 0, 1`	str (symbolic)	`n`	`register_constraint_dimension` `(b, n)` if graph_builder
`1`	str (symbolic)	`b`	none
str (symbolic)	int `n ≠ 0, 1`	`n`	`register_constraint_dimension` `(a, n)` if graph_builder
str (symbolic)	`1`	`a`	none
str `a == b`	str `a == b`	`a`	none
str `a != b`	str `a != b`	`a^b`	none (`^` means `max`)

When a symbolic dimension is paired with a concrete integer n ≠ 1, the concrete value is chosen as the output dimension and the equality is stored as a constraint via register_constraint_dimension. This avoids the need to backtrack through earlier nodes when the concrete value is discovered later: downstream operations immediately see a precise integer shape.

Parameters:

sh1 – first shape (tuple of ints and/or symbolic strings)
sh2 – second shape (tuple of ints and/or symbolic strings)
graph_builder – if not None, constraints are registered on this builder whenever a symbolic dimension is equated to a concrete integer

Returns:

resulting broadcast shape

yobx.xshape.shape_type_compute.set_shape_type_attention_microsoft(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for com.microsoft.Attention.

yobx.xshape.shape_type_compute.set_shape_type_causal_conv_with_state(self: ShapeBuilder, node: NodeProto)[source]#

Sets the output shape for com.microsoft.CausalConvWithState.

Inputs: input (N, C, L), weight (C, 1, K), bias (C) (optional), past_state (N, C, K-1) (optional). Outputs: output (same shape as input), present_state (N, C, K-1) (optional).

yobx.xshape.shape_type_compute.set_shape_type_cdist(self: ShapeBuilder, node: NodeProto)[source]#

Sets the output shape for com.microsoft.CDist.

Input A has shape (N, D) and input B has shape (M, D), so the output has shape (N, M).

yobx.xshape.shape_type_compute.set_shape_type_custom(self: ShapeBuilder, node: NodeProto, exc: bool = False)[source]#: Sets the shape and type if it can.

yobx.xshape.shape_type_compute.set_shape_type_decoder_attention(self: ShapeBuilder, node: NodeProto)[source]#

Sets the output shape for com.microsoft.DecoderAttention.

Inputs (sequence-first format):

query (S, B, H) sequence_length, batch_size, hidden_size
key (T, B, H)
q_weight (H, H)
kv_weight (H, 2*H)
bias (3*H,)
key_padding_mask (optional) (B, T)
key_cache (optional) (B, num_heads, *, head_size)
value_cache``(optional) ``(B, num_heads, *, head_size)
static_kv, use_past, has_layer_state, has_key_padding_mask – bool scalars

Outputs:

output (S, B, H)
new_key_cache (optional) (B, num_heads, *, head_size)
new_value_cache``(optional) ``(B, num_heads, *, head_size)

yobx.xshape.shape_type_compute.set_shape_type_embed_layer_normalization(self: ShapeBuilder, node: NodeProto)[source]#

Sets the output shape for com.microsoft.EmbedLayerNormalization.

Inputs: input_ids [B, S], segment_ids [B, S] (optional), word_embedding [V, D], position_embedding [P, D], segment_embedding [NS, D] (optional), gamma [D], beta [D], mask [B, S] (optional), position_ids [B, S] (optional). Outputs: output [B, S, D], mask_index [B].

yobx.xshape.shape_type_compute.set_shape_type_gated_relative_position_bias(self: ShapeBuilder, node: NodeProto)[source]#

Sets the output shape for com.microsoft.GatedRelativePositionBias.

Inputs: query_layer (batch, seq_len, num_heads*head_size), query_bias, rel_pos (1, num_heads, seq_len, seq_len), weight, bias, eco_a (1, num_heads, 1, 1), [token_offset]. Output: (batch_size, num_heads, seq_len, seq_len).

yobx.xshape.shape_type_compute.set_shape_type_greedy_search(self: ShapeBuilder, node: NodeProto)[source]#

Sets the output shape for com.microsoft.GreedySearch.

Input input_ids has shape (batch_size, sequence_length) and type INT32. Input max_length is a scalar INT32. Output sequences has shape (batch_size, max_length_value) and type INT32.

yobx.xshape.shape_type_compute.set_shape_type_group_query_attention(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for com.microsoft.GroupQueryAttention.

yobx.xshape.shape_type_compute.set_shape_type_linear_attention(self: ShapeBuilder, node: NodeProto)[source]#

Sets the output shape for com.microsoft.LinearAttention.

Inputs (3-D packed format):

query (B, T, H_q * d_k)
key (B, T, H_kv * d_k)
value (B, T, H_kv * d_v)
past_state (optional) (B, H_kv, d_k, d_v)
decay (optional)
beta (optional)

Outputs:

output (B, T, H_q * d_v)
present_state (B, H_kv, d_k, d_v)

yobx.xshape.shape_type_compute.set_shape_type_moe(self: ShapeBuilder, node: NodeProto)[source]#

Sets the output shape for com.microsoft.MoE (Mixture of Experts).

Input input has shape (num_tokens, hidden_size) or (batch_size, seq_len, hidden_size). Output output has the same shape and dtype as the input.

yobx.xshape.shape_type_compute.set_shape_type_murmur_hash3(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for com.microsoft.MurmurHash3.

yobx.xshape.shape_type_compute.set_shape_type_op_any(self: ShapeBuilder, node: NodeProto, exc: bool = False)[source]#: Sets the shape and type if it can.

yobx.xshape.shape_type_compute.set_shape_type_op_any_arg_max_min(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for ArgMax and ArgMin.

yobx.xshape.shape_type_compute.set_shape_type_op_any_attention(g: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Attention.

yobx.xshape.shape_type_compute.set_shape_type_op_any_batch_normalization(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type BatchNormalization.

yobx.xshape.shape_type_compute.set_shape_type_op_any_cast(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Cast.

yobx.xshape.shape_type_compute.set_shape_type_op_any_castlike(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type CastLike.

yobx.xshape.shape_type_compute.set_shape_type_op_any_compress(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Compress.

yobx.xshape.shape_type_compute.set_shape_type_op_any_concat(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Concat.

yobx.xshape.shape_type_compute.set_shape_type_op_any_constantofshape(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output type for ConstantOfShape.

yobx.xshape.shape_type_compute.set_shape_type_op_any_conv_max_pool(self: ShapeBuilder, node: NodeProto)[source]#

Sets the output shape for node types Conv, MaxPool.

This function defines the following functions:

conv_f1(d,s,stride) = s - (stride if d % stride == 0 else d % stride) // 2
conv_f2(d,s,stride) = (
    s - (stride if d % stride == 0 else d % stride)) // 2 + stride % 2
)
conv_f3(d,s,stride,ceil_mode,p) = ... (see the code)

yobx.xshape.shape_type_compute.set_shape_type_op_any_depth_to_space(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for DepthToSpace.

yobx.xshape.shape_type_compute.set_shape_type_op_any_dropout(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Dropout.

yobx.xshape.shape_type_compute.set_shape_type_op_any_einsum(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Einsum.

yobx.xshape.shape_type_compute.set_shape_type_op_any_expand(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Reshape.

yobx.xshape.shape_type_compute.set_shape_type_op_any_eyelike(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for EyeLike.

yobx.xshape.shape_type_compute.set_shape_type_op_any_flatten(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for Flatten.

yobx.xshape.shape_type_compute.set_shape_type_op_any_gather(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Gather.

yobx.xshape.shape_type_compute.set_shape_type_op_any_gather_elements(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type GatherElements.

yobx.xshape.shape_type_compute.set_shape_type_op_any_gathernd(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output type for GatherND.

yobx.xshape.shape_type_compute.set_shape_type_op_any_gemm(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Gemm.

yobx.xshape.shape_type_compute.set_shape_type_op_any_gemma_rotary_embedding(self: ShapeBuilder, node: NodeProto)[source]#

Sets the output shapes for node type GemmaRotaryEmbedding.

Inputs: emb (batch, seq, dim), q (batch, heads_q, seq, dim), q_rot, k (batch, heads_k, seq, dim), k_rot. Outputs: q_embed matching q, k_embed matching k.

yobx.xshape.shape_type_compute.set_shape_type_op_any_global_pool(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for GlobalAveragePool and GlobalMaxPool.

yobx.xshape.shape_type_compute.set_shape_type_op_any_gridsample(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for GridSample.

yobx.xshape.shape_type_compute.set_shape_type_op_any_identity(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output type and shape for Identity (passthrough).

yobx.xshape.shape_type_compute.set_shape_type_op_any_instance_normalization(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type InstanceNormalization.

yobx.xshape.shape_type_compute.set_shape_type_op_any_layer_normalization(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type LayerNormalization.

yobx.xshape.shape_type_compute.set_shape_type_op_any_loop(self: ShapeBuilder, node: NodeProto)[source]#

Sets output types and shapes for the Loop operator from its body subgraph.

The Loop body subgraph has the following output order: [cond_out, v_out_0, ..., v_out_K-1, scan_0, ..., scan_S-1]

The Loop node outputs are: [v_final_0, ..., v_final_K-1, scan_0_stacked, ..., scan_S-1_stacked]

Loop-carried output v_final_i has the same type and shape as body.output[i+1]. Each scan output has the same element type as the corresponding body scan output but with an extra leading dimension (the number of iterations, which is usually unknown at compile time).

yobx.xshape.shape_type_compute.set_shape_type_op_any_lp_normalization(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type LpNormalization.

yobx.xshape.shape_type_compute.set_shape_type_op_any_matmul(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type MatMul.

yobx.xshape.shape_type_compute.set_shape_type_op_any_mel_weight_matrix(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output type for MelWeightMatrix.

yobx.xshape.shape_type_compute.set_shape_type_op_any_nms(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output type for NonMaxSuppression (output is INT64).

yobx.xshape.shape_type_compute.set_shape_type_op_any_non_zero(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type NonZero.

yobx.xshape.shape_type_compute.set_shape_type_op_any_onehot(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output type for OneHot (type from values input).

yobx.xshape.shape_type_compute.set_shape_type_op_any_pad(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Pad.

yobx.xshape.shape_type_compute.set_shape_type_op_any_qlinear_matmul(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output type for QLinearMatMul.

yobx.xshape.shape_type_compute.set_shape_type_op_any_range(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for for node type Range.

yobx.xshape.shape_type_compute.set_shape_type_op_any_reduce(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for Reduce node type.

yobx.xshape.shape_type_compute.set_shape_type_op_any_reshape(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Reshape.

yobx.xshape.shape_type_compute.set_shape_type_op_any_resize(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Resize.

yobx.xshape.shape_type_compute.set_shape_type_op_any_rnn(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output types for RNN/GRU/LSTM.

yobx.xshape.shape_type_compute.set_shape_type_op_any_rotary_embedding(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type RotaryEmbedding.

yobx.xshape.shape_type_compute.set_shape_type_op_any_scatternd(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type ScatterND.

yobx.xshape.shape_type_compute.set_shape_type_op_any_shape(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output type and rank for Shape op.

yobx.xshape.shape_type_compute.set_shape_type_op_any_sign(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Sign.

yobx.xshape.shape_type_compute.set_shape_type_op_any_size(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Sign.

yobx.xshape.shape_type_compute.set_shape_type_op_any_slice(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Slice.

yobx.xshape.shape_type_compute.set_shape_type_op_any_softmax_cross_entropy_loss(self: ShapeBuilder, node: NodeProto) → bool[source]#

Sets shape/type for SoftmaxCrossEntropyLoss.

Output 0 shape depends on the reduction attribute:

"none": batch shape = input[0] shape without the class dimension (dim 1), i.e. (N,) for 2-D input or (N, d1, ..., dk) for higher-rank input.
"mean" / "sum": scalar ().

Output 1 (log_prob, optional): same shape as input[0].

yobx.xshape.shape_type_compute.set_shape_type_op_any_space_to_depth(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for SpaceToDepth.

yobx.xshape.shape_type_compute.set_shape_type_op_any_split(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Split.

yobx.xshape.shape_type_compute.set_shape_type_op_any_squeeze(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Squeeze.

yobx.xshape.shape_type_compute.set_shape_type_op_any_tile(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Tile.

yobx.xshape.shape_type_compute.set_shape_type_op_any_topk(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type TopK.

yobx.xshape.shape_type_compute.set_shape_type_op_any_transpose(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Transpose.

yobx.xshape.shape_type_compute.set_shape_type_op_any_unary(self: ShapeBuilder, node: NodeProto, itype: int | None = None)[source]#: Sets the output shape for any unary type.

yobx.xshape.shape_type_compute.set_shape_type_op_any_unsqueeze(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Unsqueeze.

yobx.xshape.shape_type_compute.set_shape_type_op_any_where(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Where.

yobx.xshape.shape_type_compute.set_shape_type_op_any_window(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for BlackmanWindow, HannWindow, HammingWindow.

yobx.xshape.shape_type_compute.set_shape_type_packed_multi_head_attention(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for com.microsoft.PackedMultiHeadAttention.

yobx.xshape.shape_type_compute.set_shape_type_relative_position_bias(self: ShapeBuilder, node: NodeProto)[source]#

Sets the output shape for com.microsoft.RelativePositionBias.

Inputs: bias_table (num_heads, num_buckets), query_length (), key_length (). Output: (1, num_heads, query_length, key_length).

yobx.xshape.shape_type_compute.set_type_shape_bias_split_gelu(g: ShapeBuilder, node: NodeProto)[source]#

Sets the shape and type for com.microsoft.BiasSplitGelu.

The operator computes Y = left * Gelu(right) after adding a bias and splitting the last dimension into two equal halves, so the output shape equals the input shape with the last dimension halved.

yobx.xshape.shape_type_compute.set_type_shape_binary_op(g: ShapeBuilder, name: str, *input_names: List[str], begin: int = 0, cmp_op: bool = False, itype: int | None = None) → bool[source]#: Sets the shape and type for a binary operator (add, mul, …).

yobx.xshape.shape_type_compute.set_type_shape_complex_module(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type ComplexModule (extracts real/imaginary part).

yobx.xshape.shape_type_compute.set_type_shape_fused_matmul(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type FusedMatMul.

yobx.xshape.shape_type_compute.set_type_shape_gemm(g: ShapeBuilder, name: str, x: str, y: str, transA: int, transB: int)[source]#: Sets the output shape for node type Gemm.

yobx.xshape.shape_type_compute.set_type_shape_matmul(g: ShapeBuilder, name: str, x: str, y: str) → bool[source]#: Sets the output shape for node type MatMul.

yobx.xshape.shape_type_compute.set_type_shape_multi_head_attention(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type MultiHeadAttention.

yobx.xshape.shape_type_compute.set_type_shape_reduce_op(g: ShapeBuilder, name: str, x: str, keepdim: int, axes: Tuple[int] | None = None)[source]#: Sets the output shape for any Reduce type.

yobx.xshape.shape_type_compute.set_type_shape_reshape(g: ShapeBuilder, name: str, input_name: str, new_shape: Sequence[int])[source]#: Sets the output shape for node type Reshape

yobx.xshape.shape_type_compute.set_type_shape_scatter_nd_of_shape(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node types ScatterNDOfShape and MaskedScatterNDOfShape.

yobx.xshape.shape_type_compute.set_type_shape_shared_input(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shapes for nodes with two outputs sharing the same inputs.

yobx.xshape.shape_type_compute.set_type_shape_to_complex(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type ToComplex (converts float to complex).

yobx.xshape.shape_type_compute.set_type_shape_transpose_2d_cast_fp16(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Transpose2DCastFP16 (transposes and casts to float16).

yobx.xshape.shape_type_compute.set_type_shape_transpose_2d_cast_fp32(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type Transpose2DCastFP32 (transposes and casts to float32).

yobx.xshape.shape_type_compute.set_type_shape_tree_ensemble(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node types TreeEnsemble and TreeEnsembleRegressor.

yobx.xshape.shape_type_compute.set_type_shape_tri_matrix(self: ShapeBuilder, node: NodeProto)[source]#: Sets the output shape for node type TriMatrix.

yobx.xshape.shape_type_compute.set_type_shape_unary_op(g: ShapeBuilder, name: str, input_name: str, itype: int | None = None) → bool[source]#: Sets the shape and type for an unary operator (abs, exp, …).

yobx.xshape.shape_type_compute.set_type_shape_unary_op_abs(g: ShapeBuilder, name: str, input_name: str, itype: int | None = None) → bool[source]#: Sets the shape and type for an unary operator (abs, exp, …).

yobx.xshape.shape_type_compute.supported_ops_in_set_shape_type_custom() → Dict[str, FrozenSet[str]][source]#

Returns the ops supported by set_shape_type_custom() grouped by domain.

Returns a dictionary mapping each ONNX domain name to a frozenset of op type names for which set_shape_type_custom() provides shape and type inference.

The special key "" (empty string) groups ops that are handled regardless of their domain (i.e. no domain check is performed for them). Local functions registered at runtime are not included because they are determined dynamically.

Returns:: Dictionary mapping domain name to a frozenset of supported op types.