Module Nir.NodeSource

Nodes descriptions

A node is composed of

descr describes several operations. When an operation shares the same name as an ONNX operation, it follows the standard defined in the ONNX IR v8 and ONNX Opset v13 standards, described here: https://onnx.ai/onnx/operators/index.html.

Nodes only require their inputs: it is assumed that a node only returns one value.

Sourcetype ty =
  1. | BFloat16
  2. | Float16
  3. | Float
  4. | UInt8
  5. | Int8
  6. | Int32
  7. | Int64
Sourcetype descr =
  1. | Constant of {
    1. data : Gentensor.t;
    }
    (*

    A constant tensor, used to store non-varying parameters during inference.

    *)
  2. | Add of {
    1. input1 : t;
    2. input2 : t;
    }
  3. | Sub of {
    1. input1 : t;
    2. input2 : t;
    }
  4. | Mul of {
    1. input1 : t;
    2. input2 : t;
    }
  5. | Div of {
    1. input1 : t;
    2. input2 : t;
    }
  6. | Matmul of {
    1. input1 : t;
    2. input2 : t;
    }
  7. | QLinearMatMul of {
    1. inputA : t;
    2. inputA_scale : t;
    3. inputA_zero_point : t;
    4. inputB : t;
    5. inputB_scale : t;
    6. inputB_zero_point : t;
    7. y_scale : t;
    8. y_zero_point : t;
    }
  8. | Gemm of {
    1. inputA : t;
    2. inputB : t;
    3. inputC : t Base.option;
    4. alpha : Base.float;
    5. beta : Base.float;
    6. transA : Base.int;
    7. transB : Base.int;
    }
  9. | QGemm of {
    1. inputA : t;
    2. inputA_scale : t;
    3. inputA_zero_point : t;
    4. inputB : t;
    5. inputB_scale : t;
    6. inputB_zero_point : t;
    7. inputC : t Base.option;
    8. y_scale : t Base.option;
    9. y_zero_point : t Base.option;
    10. alpha : Base.float;
    11. transA : Base.int;
    12. transB : Base.int;
    }
    (*

    Not an ONNX operator of the default domain. Documentation at:

    • https://xadupre.github.io/draft/inference/operators/onnx_commicrosoft_QGemm.html
    • https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#commicrosoftqgemm
    *)
  10. | LogSoftmax
  11. | ReLu of {
    1. input : t;
    }
  12. | Transpose of {
    1. input : t;
      (*

      Called "data" in ONNX documentation: https://onnx.ai/onnx/operators/onnx__Transpose.html .

      *)
    2. perm : Base.int Base.list;
    }
  13. | Squeeze of {
    1. data : t;
    2. axes : t Base.option;
    }
  14. | MaxPool
  15. | Conv
  16. | Reshape of {
    1. input : t;
    2. shape : t;
    }
  17. | Flatten of {
    1. input : t;
    2. axis : Base.int;
    }
  18. | Identity of {
    1. input : t;
    }
  19. | Input of {
    1. shape : Shape.t;
    }
  20. | RW_Linearized_ReLu
  21. | Concat of {
    1. inputs : t Base.list;
    2. axis : Base.int;
    }
  22. | Gather of {
    1. input : t;
    2. indices : t;
    3. axis : Base.int;
    }
  23. | ReduceSum of {
    1. input : t;
    2. axes : t Base.option;
    3. keepdims : Base.int;
    4. noop_with_empty_axes : Base.int;
    }
  24. | GatherND of {
    1. data : t;
    2. indices : t;
    3. batch_dims : Base.int;
    }
  25. | RandomNormal of {
    1. dtype : Base.int;
    2. mean : Base.float;
    3. scale : Base.float;
    4. seed : Base.float;
    5. shape : Base.int Base.array;
    }
  26. | Abs of {
    1. input : t;
    }
  27. | Log of {
    1. input : t;
    }
  28. | Sign of {
    1. input : t;
    }
  29. | ArgMax of {
    1. input : t;
    2. axis : Base.int;
    3. keepdims : Base.bool;
    }
  30. | Pow of {
    1. input1 : t;
    2. input2 : t;
    }
  31. | QuantizeLinear of {
    1. x : t;
    2. y_scale : t;
    3. y_zero_point : t Base.option;
    4. axis : Base.int;
    }
  32. | DequantizeLinear of {
    1. x : t;
    2. x_scale : t;
    3. x_zero_point : t Base.option;
    4. axis : Base.int;
    }
Sourceand t = private {
  1. id : Base.int;
  2. descr : descr;
  3. shape : Shape.t;
  4. ty : ty;
    (*

    Describes the shape of the result of the node computation.

    *)
}
Sourceval equal : t -> t -> Base.bool
include Base.Hashtbl.Key.S with type t := t
Sourceval compare : t -> t -> int
Sourceval sexp_of_t : t -> Sexplib0.Sexp.t
Sourceval hash : t -> int

Two ts that compare equal must have equal hashes for the hashtable to behave properly.

include Base.Comparator.S with type t := t
Sourcetype comparator_witness
Sourceval create : descr -> t

create descr returns a value of type node with proper indexing and the shape according to the ONNX semantic.

Sourceval gather_int : ?encode:Base.bool -> t -> Base.int -> t

gather_int n i gathers the ith element of the node n.

Sourceval gather_ints : ?encode:Base.bool -> t -> Base.int Base.list -> t

gather_ints n l gathers the ith element of the node n for each element i of l.

Sourceval map : (t -> t) -> t -> t

map f n replace the direct inputs i of n by f i

Sourceval map_rec : (t -> t) -> t -> t

map_rec f n replace top-bottom the nodes i accessible from n by f i

Sourceval replace_input : (Base.unit -> t) -> t -> t

replace_input f n replace the input in n by f ()

Sourceval preds : t -> t Base.list

Direct predecessors of a t.

Sourceval iter_rec : (t -> Base.unit) -> t -> Base.unit

Iterate on the predecessors of a t and itself. Repect topological order.

Sourceval compute_shape : t -> Shape.t
Sourceval (+) : t -> t -> t
Sourceval (*) : t -> t -> t
Sourceval (-) : t -> t -> t
Sourceval mul_float : t -> Base.float -> t
Sourceval div_float : ?encode:Base.bool -> t -> Base.float -> t
Sourceval concat_0 : t Base.list -> t
Sourceval reshape : Shape.t -> t -> t
Sourceval sum_list : ?shp:Shape.t -> t Base.list -> t

sum_list shp ns is a node corresponding to the sum of the nodes in ns. If ns is empty, this returns a tensor of shape shp filled with 0s. By default, shp is a single float.

Sourceval partial_dot_product : ?shp:Shape.t -> t Base.array -> t Base.array -> Base.int -> Base.int -> t

partial_dot_product shp arr1 arr2 first last where arr1 = [\|n11, n12, ..., n1k1\|] and arr2 = [\|n21, n22, ..., n2k2\|] is a node corresponding to (n1first * n2first) + (n1first + 1 * n2first + 1) + ... + (n1last - 1 * n2last - 1) if this exists. It is assumed that arr1 and arr2 contain tensors with same shape. Edge cases include:

  • if last > length n1 or last > length n2, then fails
  • if last >= first, then returns a tensor where all values are initialized to 0. The shape of this tensor is determined using the following order:
  1. if length arr1 <> 0 then use the shape of arr1.(0)
  2. if length arr2 <> 0 then use the shape of arr2.(0)
  3. if shp <> None, then use shp
  4. otherwise, fails

transpose perm id is the position of the component at position id when the permutation perm is applied as the result of a tensor transposition. For instance, if id = [\|10; 20; 30\|] and \|perm = [2;0;1]\|, then transpose perm id will equal [\|30; 10; 20\|]. The empty permutation is represented by [] and is interpreted as [r-1; r-2; ...; 1; 0] (following <https://onnx.ai/onnx/operators/onnx__Transpose.html#transpose-23>); otherwise it is assumed that perm is a permutation of [0; 1; ...; r-1].

untranspose is the reverse of transpose so that untranspose perm @@ transpose perm id equals id.

flatten shp axis id computes the position of the component at position id when the flattening on axis is performed over shape shp. Following the definition of Flatten <https://onnx.ai/onnx/operators/onnx__Flatten.html#flatten-23>, axis can be negative in which case it computes from the end. For instance, if the shape is 10 * 10 * 10 * 10 * 10 (rank 5), axis = 3, and id = [\|1; 2; 3; 4; 5\|], then the result will be [\|123; 45\|] (computed as [\| 3 + (10 * ( 2 + 10 * 1)); 5 + 10 * 4) \|]).

unflatten is the reverse of flatten so that unflatten sh axis @@ flatten sh axis id equals id (for correct inputs id).

Sourceval encode_qgemm : descr -> descr

Encode QGemm operator in terms of DequantizeLinear, QuantizeLinear and Gemm operators.