yobx.xoptim.patterns.onnx_gather#

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

Simplifies Gather(Concat(..., X, ..., axis=0), cst1) into Gather(X, cst2) where X is a 1D tensor and cst1, cst2 are 0D or 1D integer tensors.

This applies when:

  • The Concat axis is 0.

  • Exactly one Concat input (X) is not a constant; all others are constants with known sizes.

  • cst1 is a constant 0D or 1D INT64 tensor with non-negative values.

  • All indices in cst1 fall within X’s slice of the concatenated tensor.

The adjusted index is cst2 = cst1 - offset where offset is the sum of the sizes of all constant inputs that precede X in the Concat input list.

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 INT64(n)"])
    I_C1(["C1 INT64(3)"])
    I_C2(["C2 INT64(2)"])
    I_cst1(["cst1 INT64()"])

    Concat_0[["Concat(., ., ., axis=0)"]]
    Gather_1[["Gather(., .)"]]

    I_C1 -->|"INT64(3)"| Concat_0
    I_X -->|"INT64(n)"| Concat_0
    I_C2 -->|"INT64(2)"| Concat_0
    Concat_0 -->|"INT64(3+n+2)"| Gather_1
    I_cst1 -->|"INT64()"| Gather_1

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

    class I_X,I_C1,I_C2,I_cst1,O_Y ioNode
    class Concat_0,Gather_1 opNode

Outcome of the fusion:

        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 INT64(n)"])
    I_cst2(["cst2 INT64()"])

    Gather_0[["Gather(., .)"]]

    I_X -->|"INT64(n)"| Gather_0
    I_cst2 -->|"INT64()"| Gather_0

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

    class I_X,I_cst2,O_Y ioNode
    class Gather_0 opNode
apply(g: GraphBuilder, concat_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_gather.GatherGatherPattern(verbose: int = 0, priority: int = 0)[source]#

Simplifies two consecutive Gather operations into one.

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(n, a, b)"])
    I_cst1(["cst1 INT64(k)"])
    I_cst2(["cst2 INT64()"])

    Constant_0[["Constant() -#gt; cst1"]]
    Constant_1[["Constant() -#gt; cst2"]]
    Gather_2[["Gather(., ., axis=0)"]]
    Gather_3[["Gather(., ., axis=0)"]]

    I_X -->|"FLOAT(n, a, b)"| Gather_2
    Constant_0 -->|"INT64(k)"| Gather_2
    Gather_2 -->|"FLOAT(k, a, b)"| Gather_3
    Constant_1 -->|"INT64()"| Gather_3

    O_Y(["Y FLOAT(a, b)"])
    Gather_3 --> O_Y

    class I_X,I_cst1,I_cst2,O_Y ioNode
    class Constant_0,Constant_1 constNode
    class Gather_2,Gather_3 opNode

Outcome of the fusion:

        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(n, a, b)"])
    I_cst3(["cst3 INT64()"])

    Constant_0[["Constant() -#gt; cst3"]]
    Gather_1[["Gather(., ., axis=0)"]]

    I_X -->|"FLOAT(n, a, b)"| Gather_1
    Constant_0 -->|"INT64()"| Gather_1

    O_Y(["Y FLOAT(a, b)"])
    Gather_1 --> O_Y

    class I_X,I_cst3,O_Y ioNode
    class Constant_0 constNode
    class Gather_1 opNode

The composed index is cst3 = cst1[cst2]. This applies when both Gather operations use axis=0, the inner Gather indices cst1 are a 1-D constant array, and the outer Gather indices cst2 are a constant of any shape (scalar or 1-D).

apply(g: GraphBuilder, inner_node: NodeProto, outer_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.