yobx.tensorflow.ops.shape_ops#

Converters for TF shape-manipulation ops: Reshape, Squeeze, Shape, StridedSlice, Pack, ExpandDims, Transpose, Cast.

Reshape / squeeze#

Reshape, Squeeze

Type casting#

Cast

Shape / indexing#

Shape, StridedSlice

Stack / pack#

Pack

Dimension insertion / permutation#

ExpandDims, Transpose

yobx.tensorflow.ops.shape_ops.convert_cast(g: GraphBuilderExtendedProtocol, sts: Dict[str, Any], outputs: List[str], op: Operation) → str[source]#

Converts TF Cast → ONNX Cast.

The target dtype is read from the TF op’s DstT attribute and mapped to the corresponding ONNX TensorProto integer type.

yobx.tensorflow.ops.shape_ops.convert_expand_dims(g: GraphBuilderExtendedProtocol, sts: Dict[str, Any], outputs: List[str], op: Operation) → str[source]#

Converts TF ExpandDims (tf.expand_dims) → ONNX Unsqueeze.

The dim input (second input) must be a scalar constant; its value is read and forwarded as the axes argument of Unsqueeze.

yobx.tensorflow.ops.shape_ops.convert_pack(g: GraphBuilderExtendedProtocol, sts: Dict[str, Any], outputs: List[str], op: Operation) → str[source]#

Converts TF Pack (tf.stack) → ONNX Unsqueeze + Concat.

Pack stacks N tensors along a new axis. In ONNX this is expressed as unsqueezing each input along the new axis and then concatenating them.

Each input is cast to int64 so that the resulting tensor can be used as an ONNX Reshape shape argument (which requires int64).

yobx.tensorflow.ops.shape_ops.convert_reshape(g: GraphBuilderExtendedProtocol, sts: Dict[str, Any], outputs: List[str], op: Operation) → str[source]#

Converts TF Reshape → ONNX Reshape.

The TF shape input (second input) may be int32; it is cast to int64 as required by the ONNX Reshape specification.

When ONNX shape inference cannot determine the output shape (e.g. because the shape tensor contains a dynamic batch dimension), the TF-inferred output shape is used as a fallback so that downstream shape-inference assertions are satisfied.

yobx.tensorflow.ops.shape_ops.convert_shape(g: GraphBuilderExtendedProtocol, sts: Dict[str, Any], outputs: List[str], op: Operation) → str[source]#

Converts TF Shape → ONNX Shape.

TF’s Shape op may output int32 (via the out_type attribute), but the ONNX Shape operator always returns int64. The output is kept as int64 throughout the ONNX graph; downstream converters (e.g. StridedSlice, Pack) cast their non-int64 constant inputs to match.

yobx.tensorflow.ops.shape_ops.convert_squeeze(g: GraphBuilderExtendedProtocol, sts: Dict[str, Any], outputs: List[str], op: Operation) → str[source]#

Converts TF Squeeze → ONNX Squeeze.

When squeeze_dims is empty all size-1 dimensions are removed (matching TF’s default behaviour). When specific axes are given they are passed as an int64 tensor input to the ONNX node (opset ≥ 13 API).

yobx.tensorflow.ops.shape_ops.convert_strided_slice(g: GraphBuilderExtendedProtocol, sts: Dict[str, Any], outputs: List[str], op: Operation) → str[source]#

Converts TF StridedSlice → ONNX Gather or Slice (+ optional Squeeze).

Only the common cases are supported:

  • No begin_mask, end_mask, ellipsis_mask, or new_axis_mask bits set.

  • shrink_axis_mask is supported by using Gather with a scalar index for the typical single-element-extraction pattern (e.g. as emitted by tf.keras.layers.Flatten).

This covers the pattern generated by tf.keras.layers.Flatten and similar common usages of StridedSlice on shape tensors.

yobx.tensorflow.ops.shape_ops.convert_transpose(g: GraphBuilderExtendedProtocol, sts: Dict[str, Any], outputs: List[str], op: Operation) → str[source]#

Converts TF Transpose → ONNX Transpose.

The permutation is given as a second input tensor (a 1-D int32/int64 constant); its value is read and forwarded as the perm attribute.