package caisar

  1. Overview
  2. Docs

Module Nir.NodeSource

Nodes descriptions

A node is composed of

  • a unique id of type int
  • a node description of type descr

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. | Float
  2. | 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. | 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;
    }
  8. | LogSoftmax
  9. | ReLu of {
    1. input : t;
    }
  10. | Transpose of {
    1. input : t;
      (*

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

      *)
    2. perm : Base.int Base.list;
    }
  11. | Squeeze of {
    1. data : t;
    2. axes : t Base.option;
    }
  12. | MaxPool
  13. | Conv
  14. | Reshape of {
    1. input : t;
    2. shape : t;
    }
  15. | Flatten of {
    1. input : t;
    2. axis : Base.int;
    }
  16. | Identity of {
    1. input : t;
    }
  17. | Input of {
    1. shape : Shape.t;
    }
  18. | RW_Linearized_ReLu
  19. | Concat of {
    1. inputs : t Base.list;
    2. axis : Base.int;
    }
  20. | Gather of {
    1. input : t;
    2. indices : t;
    3. axis : Base.int;
    }
  21. | ReduceSum of {
    1. input : t;
    2. axes : t Base.option;
    3. keepdims : Base.int;
    4. noop_with_empty_axes : Base.int;
    }
  22. | GatherND of {
    1. data : t;
    2. indices : t;
    3. batch_dims : Base.int;
    }
  23. | RandomNormal of {
    1. dtype : Base.int;
    2. mean : Base.float;
    3. scale : Base.float;
    4. seed : Base.float;
    5. shape : Base.int Base.array;
    }
  24. | Abs of {
    1. input : t;
    }
  25. | Log of {
    1. input : t;
    }
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
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 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
OCaml

Innovation. Community. Security.