yobx.xoptim.patterns.onnx_shape#

class yobx.xoptim.patterns.onnx_shape.GatherShapePattern(verbose: int = 0, priority: int = 0)[source]#

Simplifies Gather(Shape(X), indices) into Shape(X, start=s, end=e) when indices is a constant 1-D int64 array that forms a contiguous ascending range [s, s+1, ..., e-1].

This avoids materialising the full shape vector only to slice it immediately afterwards. The Shape node may already carry start / end attributes (ONNX opset ≥ 15); those are taken into account when computing the absolute indices in X’s dimension space.

Model with nodes to be fused:

        graph TD

    classDef ioNode fill:#dfd,stroke:#333,color:#333
    classDef initNode fill:#cccc00,stroke:#333,color:#333
    classDef constNode fill:#f9f,stroke:#333,stroke-width:2px,color:#333
    classDef opNode fill:#bbf,stroke:#333,stroke-width:2px,color:#333

    I_X(["X FLOAT(a, b, c, d)"])
    C_idx(["idx INT64[3]"])

    Shape_0[["Shape(.)"]]
    Gather_1[["Gather(., [0, 1, 2])"]]

    I_X -->|"FLOAT(a, b, c, d)"| Shape_0
    C_idx -->|"INT64(3)"| Gather_1
    Shape_0 -->|"INT64(4)"| Gather_1

    O_Y(["Y INT64(3)"])
    Gather_1 --> O_Y

    class I_X,O_Y ioNode
    class C_idx constNode
    class Shape_0,Gather_1 opNode

Outcome of the fusion:

        graph TD

    classDef ioNode fill:#dfd,stroke:#333,color:#333
    classDef opNode fill:#bbf,stroke:#333,stroke-width:2px,color:#333

    I_X(["X FLOAT(a, b, c, d)"])

    Shape_0[["Shape(., start=0, end=3)"]]

    I_X -->|"FLOAT(a, b, c, d)"| Shape_0

    O_Y(["Y INT64(3)"])
    Shape_0 --> O_Y

    class I_X,O_Y ioNode
    class Shape_0 opNode
apply(g: GraphBuilderPatternOptimization, shape_node: NodeProto, gather_node: NodeProto) → List[NodeProto][source]#

The method does the rewriting. It assumes it can happen. It takes a list of nodes impacted by the rewriting assumes no other pattern optimizer will be modify them. It receives the list of nodes returned by method apply. Since it is a list of argument, method match can include None values. The method returns the new nodes. The optimizer considers that any node given to this function is removed from the graph, and any node returned by it are added. If a received node must be kept, it must be added to the list of returned node.

Parameters:

nodes – nodes returned by method match, there are then removed

Returns:

nodes to add to graph.

match(g: GraphBuilderPatternOptimization, node: NodeProto, matched: List[MatchResult]) → MatchResult | None[source]#

Determines nodes around node which can be rewritten.

Parameters:

  • g – is a GraphBuilderPatternOptimization, it holds all the existing nodes, is able to return any information about type, shape, the node before, the node after another one.

  • node – the matching must determine if some nodes around this one are part of set of nodes this pattern optimizer can rewrite. From there, the function explores wherever it needs, checking any condition it needs.

  • matched – usually unused, it returns of nodes already matching a pattern

The method must not modify the graph. The method returns None if no match is found or an instance of class MatchResult. It must contain:

  • a list of nodes involved in the rewriting. It does not mean all of them will be removed but all of them are needed to do the rewriting and must not be impacted by other pattern optimizer.

  • A function doing the rewriting (usually method apply of the pattern class).

  • An existing node where the rewritten nodes can be inserted. Knowing it makes it faster to rewriter. If not specified, the optimizer will automatically determine the position of the new nodes.

class yobx.xoptim.patterns.onnx_shape.ShapeBasedShapeShapeAddPattern(verbose: int = 0, priority: int = 0)[source]#

Tries to find another way to get a dimension obtained with the addition of two.

apply(g: GraphBuilder, shape1_node: NodeProto, shape2_node: NodeProto, add_node: NodeProto) → List[NodeProto][source]#

The method does the rewriting. It assumes it can happen. It takes a list of nodes impacted by the rewriting assumes no other pattern optimizer will be modify them. It receives the list of nodes returned by method apply. Since it is a list of argument, method match can include None values. The method returns the new nodes. The optimizer considers that any node given to this function is removed from the graph, and any node returned by it are added. If a received node must be kept, it must be added to the list of returned node.

Parameters:

nodes – nodes returned by method match, there are then removed

Returns:

nodes to add to graph.

match(g: GraphBuilderPatternOptimization, node: NodeProto, matched: List[MatchResult]) → MatchResult | None[source]#

Determines nodes around node which can be rewritten.

Parameters:

  • g – is a GraphBuilderPatternOptimization, it holds all the existing nodes, is able to return any information about type, shape, the node before, the node after another one.

  • node – the matching must determine if some nodes around this one are part of set of nodes this pattern optimizer can rewrite. From there, the function explores wherever it needs, checking any condition it needs.

  • matched – usually unused, it returns of nodes already matching a pattern

The method must not modify the graph. The method returns None if no match is found or an instance of class MatchResult. It must contain:

  • a list of nodes involved in the rewriting. It does not mean all of them will be removed but all of them are needed to do the rewriting and must not be impacted by other pattern optimizer.

  • A function doing the rewriting (usually method apply of the pattern class).

  • An existing node where the rewritten nodes can be inserted. Knowing it makes it faster to rewriter. If not specified, the optimizer will automatically determine the position of the new nodes.

class yobx.xoptim.patterns.onnx_shape.ShapeTransposePattern(verbose: int = 0, priority: int = 0)[source]#

Replaces Shape(Transpose(X, perm)) by Gather(Shape(X), perm_indices) so that the expensive Transpose on the full data tensor is avoided.

The key observation is that the shape of Transpose(X, perm) is simply a permuted view of the shape of X. The permutation indices are known at optimisation time (they are an attribute of the Transpose node), so we can extract the desired dimensions directly from Shape(X) using a Gather with the (sub-)permutation as the index tensor.

For X of shape (a, b, c) and perm=[2, 0, 1] the transformation is:

# Before
xt = Transpose(X, perm=[2, 0, 1])   # (c, a, b)
Y  = Shape(xt)                       # [c, a, b]

# After
sx   = Shape(X)                      # [a, b, c]
perm = Initializer([2, 0, 1])
Y    = Gather(sx, perm, axis=0)      # [c, a, b]

Shape’s optional start/end attributes are respected: the permutation slice perm[start:end] is used as the Gather indices.

Model with nodes to be fused:

        graph TD

    classDef ioNode fill:#dfd,stroke:#333,color:#333
    classDef initNode fill:#cccc00,stroke:#333,color:#333
    classDef constNode fill:#f9f,stroke:#333,stroke-width:2px,color:#333
    classDef opNode fill:#bbf,stroke:#333,stroke-width:2px,color:#333

    I_X(["X FLOAT(a, b, c)"])

    Transpose_0[["Transpose(., perm=[2, 0, 1])"]]
    Shape_1[["Shape(.)"]]

    I_X -->|"FLOAT(a, b, c)"| Transpose_0
    Transpose_0 -->|"FLOAT(c, a, b)"| Shape_1

    O_Y(["Y INT64(3)"])
    Shape_1 --> O_Y

    class I_X,O_Y ioNode
    class Transpose_0,Shape_1 opNode

Outcome of the fusion:

        graph TD

    classDef ioNode fill:#dfd,stroke:#333,color:#333
    classDef constNode fill:#f9f,stroke:#333,stroke-width:2px,color:#333
    classDef opNode fill:#bbf,stroke:#333,stroke-width:2px,color:#333

    I_X(["X FLOAT(a, b, c)"])
    C_perm(["perm INT64[2, 0, 1]"])

    Shape_s[["Shape(.)"]]
    Gather_0[["Gather(., ., axis=0)"]]

    I_X -->|"FLOAT(a, b, c)"| Shape_s
    Shape_s -->|"INT64(3)"| Gather_0
    C_perm -->|"INT64(3)"| Gather_0

    O_Y(["Y INT64(3)"])
    Gather_0 --> O_Y

    class I_X,O_Y ioNode
    class C_perm constNode
    class Shape_s,Gather_0 opNode
apply(g: GraphBuilderPatternOptimization, tr_node: NodeProto, shape_node: NodeProto) → List[NodeProto][source]#

The method does the rewriting. It assumes it can happen. It takes a list of nodes impacted by the rewriting assumes no other pattern optimizer will be modify them. It receives the list of nodes returned by method apply. Since it is a list of argument, method match can include None values. The method returns the new nodes. The optimizer considers that any node given to this function is removed from the graph, and any node returned by it are added. If a received node must be kept, it must be added to the list of returned node.

Parameters:

nodes – nodes returned by method match, there are then removed

Returns:

nodes to add to graph.

match(g: GraphBuilderPatternOptimization, node: NodeProto, matched: List[MatchResult]) → MatchResult | None[source]#

Determines nodes around node which can be rewritten.

Parameters:

  • g – is a GraphBuilderPatternOptimization, it holds all the existing nodes, is able to return any information about type, shape, the node before, the node after another one.

  • node – the matching must determine if some nodes around this one are part of set of nodes this pattern optimizer can rewrite. From there, the function explores wherever it needs, checking any condition it needs.

  • matched – usually unused, it returns of nodes already matching a pattern

The method must not modify the graph. The method returns None if no match is found or an instance of class MatchResult. It must contain:

  • a list of nodes involved in the rewriting. It does not mean all of them will be removed but all of them are needed to do the rewriting and must not be impacted by other pattern optimizer.

  • A function doing the rewriting (usually method apply of the pattern class).

  • An existing node where the rewritten nodes can be inserted. Knowing it makes it faster to rewriter. If not specified, the optimizer will automatically determine the position of the new nodes.

class yobx.xoptim.patterns.onnx_shape.UnsqueezeShapePattern(verbose: int = 0, priority: int = 0)[source]#

Replaces Shape(Unsqueeze(X, axes)) by a Concat of Shape(X, start=s, end=e) slices interleaved with constant [1] tensors at the inserted axis positions.

The key observation is that Shape(Unsqueeze(X, axes)) produces a shape vector that is identical to Shape(X) with 1 entries inserted at the axes positions. By splitting Shape(X) into segments and concatenating them with the constant 1 values, the Unsqueeze on the (potentially large) data tensor is avoided entirely while the output shape vector remains bit-for-bit identical.

For X of shape (a, b, c) and axes=[1] the transformation is:

# Before
xu = Unsqueeze(X, [1])        # (a, 1, b, c)
Y  = Shape(xu)                # [a, 1, b, c]

# After
s0 = Shape(X, start=0, end=1) # [a]
s1 = Shape(X, start=1, end=3) # [b, c]
Y  = Concat([s0, [1], s1])    # [a, 1, b, c]

Model with nodes to be fused:

        graph TD

    classDef ioNode fill:#dfd,stroke:#333,color:#333
    classDef initNode fill:#cccc00,stroke:#333,color:#333
    classDef constNode fill:#f9f,stroke:#333,stroke-width:2px,color:#333
    classDef opNode fill:#bbf,stroke:#333,stroke-width:2px,color:#333

    I_X(["X FLOAT(a, b, c)"])
    I_axes(["axes INT64(1)"])

    Constant_0[["Constant() -#gt; axes"]]
    Unsqueeze_1[["Unsqueeze(., .)"]]
    Shape_2[["Shape(.)"]]

    I_X -->|"FLOAT(a, b, c)"| Unsqueeze_1
    Constant_0 -->|"INT64(1)"| Unsqueeze_1
    Unsqueeze_1 -->|"FLOAT(a, 1, b, c)"| Shape_2

    O_Y(["Y INT64(4)"])
    Shape_2 --> O_Y

    class I_X,I_axes,O_Y ioNode
    class Constant_0 constNode
    class Unsqueeze_1,Shape_2 opNode

Outcome of the fusion:

        graph TD

    classDef ioNode fill:#dfd,stroke:#333,color:#333
    classDef constNode fill:#f9f,stroke:#333,stroke-width:2px,color:#333
    classDef opNode fill:#bbf,stroke:#333,stroke-width:2px,color:#333

    I_X(["X FLOAT(a, b, c)"])
    C1(["const INT64[1]"])

    Shape_s0[["Shape(., start=0, end=1)"]]
    Shape_s1[["Shape(., start=1, end=3)"]]
    Concat_2[["Concat(axis=0)"]]

    I_X -->|"FLOAT(a, b, c)"| Shape_s0
    I_X -->|"FLOAT(a, b, c)"| Shape_s1
    Shape_s0 -->|"INT64(1)"| Concat_2
    C1 -->|"INT64(1)"| Concat_2
    Shape_s1 -->|"INT64(2)"| Concat_2

    O_Y(["Y INT64(4)"])
    Concat_2 --> O_Y

    class I_X,O_Y ioNode
    class C1 constNode
    class Shape_s0,Shape_s1,Concat_2 opNode
apply(g: GraphBuilderPatternOptimization, unsq_node: NodeProto, shape_node: NodeProto) → List[NodeProto][source]#

The method does the rewriting. It assumes it can happen. It takes a list of nodes impacted by the rewriting assumes no other pattern optimizer will be modify them. It receives the list of nodes returned by method apply. Since it is a list of argument, method match can include None values. The method returns the new nodes. The optimizer considers that any node given to this function is removed from the graph, and any node returned by it are added. If a received node must be kept, it must be added to the list of returned node.

Parameters:

nodes – nodes returned by method match, there are then removed

Returns:

nodes to add to graph.

match(g: GraphBuilderPatternOptimization, node: NodeProto, matched: List[MatchResult]) → MatchResult | None[source]#

Determines nodes around node which can be rewritten.

Parameters:

  • g – is a GraphBuilderPatternOptimization, it holds all the existing nodes, is able to return any information about type, shape, the node before, the node after another one.

  • node – the matching must determine if some nodes around this one are part of set of nodes this pattern optimizer can rewrite. From there, the function explores wherever it needs, checking any condition it needs.

  • matched – usually unused, it returns of nodes already matching a pattern

The method must not modify the graph. The method returns None if no match is found or an instance of class MatchResult. It must contain:

  • a list of nodes involved in the rewriting. It does not mean all of them will be removed but all of them are needed to do the rewriting and must not be impacted by other pattern optimizer.

  • A function doing the rewriting (usually method apply of the pattern class).

  • An existing node where the rewritten nodes can be inserted. Knowing it makes it faster to rewriter. If not specified, the optimizer will automatically determine the position of the new nodes.