Source file node.ml

(**************************************************************************)
(*                                                                        *)
(*  This file is part of CAISAR.                                          *)
(*                                                                        *)
(*  Copyright (C) 2025                                                    *)
(*    CEA (Commissariat à l'énergie atomique et aux énergies              *)
(*         alternatives)                                                  *)
(*                                                                        *)
(*  You can redistribute it and/or modify it under the terms of the GNU   *)
(*  Lesser General Public License as published by the Free Software       *)
(*  Foundation, version 2.1.                                              *)
(*                                                                        *)
(*  It is distributed in the hope that it will be useful,                 *)
(*  but WITHOUT ANY WARRANTY; without even the implied warranty of        *)
(*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the          *)
(*  GNU Lesser General Public License for more details.                   *)
(*                                                                        *)
(*  See the GNU Lesser General Public License version 2.1                 *)
(*  for more details (enclosed in the file licenses/LGPLv2.1).            *)
(*                                                                        *)
(**************************************************************************)
open Base

type ty =
  | BFloat16
  | Float16
  | Float
  | UInt8
  | Int8
  | Int32
  | Int64
[@@deriving show, eq]

(** TODO: add the information needed to compute the shape *)
type descr =
  | Constant of { data : Gentensor.t }
  | Add of {
      input1 : t;
      input2 : t;
    }
  | Sub of {
      input1 : t;
      input2 : t;
    }
  | Mul of {
      input1 : t;
      input2 : t;
    }
  | Div of {
      input1 : t;
      input2 : t;
    }
  | Matmul of {
      input1 : t;
      input2 : t;
    }
  | QLinearMatMul of {
      inputA : t;
      inputA_scale : t;
      inputA_zero_point : t;
      inputB : t;
      inputB_scale : t;
      inputB_zero_point : t;
      y_scale : t;
      y_zero_point : t;
    }
  | Gemm of {
      inputA : t;
      inputB : t;
      inputC : t option;
      alpha : float;
      beta : float;
      transA : int;
      transB : int;
    }
  | QGemm of {
      inputA : t;
      inputA_scale : t;
      inputA_zero_point : t;
      inputB : t;
      inputB_scale : t;
      inputB_zero_point : t;
      inputC : t option;
      y_scale : t option;
      y_zero_point : t option;
      alpha : float;
      transA : int;
      transB : int;
    }
  | LogSoftmax
  | ReLu of { input : t }
  | Transpose of {
      input : t;
        (* called "data" in ONNX documentation :
           https://onnx.ai/onnx/operators/onnx__Transpose.html*)
      perm : int list;
    }
  | Squeeze of {
      data : t;
      axes : t option; (* data int64 *)
    }
  | MaxPool
  | Conv
  | Reshape of {
      input : t;
      shape : t; (* data int64 *)
    }
  | Flatten of {
      input : t;
      axis : int;
    }
  | Identity of { input : t }
  | Input of { shape : Shape.t }
  | RW_Linearized_ReLu
  | Concat of {
      inputs : t list;
      axis : int;
    }
  | Gather of {
      input : t;
      indices : t;
      axis : int;
    }
  | ReduceSum of {
      input : t;
      axes : t option;
      keepdims : int;
      noop_with_empty_axes : int;
    }
  | GatherND of {
      data : t;
      indices : t;
      batch_dims : int;
    }
  | RandomNormal of {
      dtype : int;
      mean : float;
      scale : float;
      seed : float;
      shape : int array;
    }
  | Abs of { input : t }
  | Log of { input : t }
  | Sign of { input : t }
  | ArgMax of {
      input : t;
      axis : int;
      keepdims : bool;
    }
  | Pow of {
      input1 : t;
      input2 : t;
    }
  | QuantizeLinear of {
      x : t;
      y_scale : t;
      y_zero_point : t option;
      axis : int;
    }
  | DequantizeLinear of {
      x : t;
      x_scale : t;
      x_zero_point : t option;
      axis : int;
    }

and t = {
  id : int;
  descr : descr; [@printer fun fmt d -> pp_descr fmt d]
  shape : Shape.t;
  ty : ty;
}

let pp_descr fmt descr =
  match descr with
  | Input { shape } -> Fmt.pf fmt "Input: %a" Shape.pp shape
  | Transpose { perm; _ } ->
    Fmt.pf fmt "Transpose: [%a]" Fmt.(list ~sep:semi int) perm
  | Constant { data = Int64 b } when Shape.size (Tensor.shape b) < 3 ->
    Fmt.pf fmt "Constant[%a]" Fmt.(list ~sep:comma int64) (Tensor.flatten b)
  | Constant _ -> Fmt.pf fmt "Constant"
  | Add _ -> Fmt.pf fmt "Add"
  | Sub _ -> Fmt.pf fmt "Sub"
  | Mul _ -> Fmt.pf fmt "Mul"
  | Div _ -> Fmt.pf fmt "Div"
  | Matmul _ -> Fmt.pf fmt "Matmul"
  | QLinearMatMul _ -> Fmt.pf fmt "QLinearMatMul"
  | Gemm _ -> Fmt.pf fmt "Gemm"
  | QGemm _ -> Fmt.pf fmt "QGemm"
  | LogSoftmax -> Fmt.pf fmt "LogSoftmax"
  | ReLu _ -> Fmt.pf fmt "ReLu"
  | Squeeze _ -> Fmt.pf fmt "Squeeze"
  | MaxPool -> Fmt.pf fmt "MaxPool"
  | Conv -> Fmt.pf fmt "Conv"
  | Reshape _ -> Fmt.pf fmt "Reshape"
  | Flatten _ -> Fmt.pf fmt "Flatten"
  | Identity _ -> Fmt.pf fmt "Identity"
  | RW_Linearized_ReLu -> Fmt.pf fmt "RW_Linearized_ReLu"
  | Concat { axis; _ } -> Fmt.pf fmt "Concat[%i]" axis
  | Gather _ -> Fmt.pf fmt "Gather"
  | ReduceSum _ -> Fmt.pf fmt "ReduceSum"
  | GatherND _ -> Fmt.pf fmt "GatherND"
  | RandomNormal _ -> Fmt.pf fmt "RandomNormal"
  | Abs _ -> Fmt.pf fmt "Abs"
  | Log _ -> Fmt.pf fmt "Log"
  | Sign _ -> Fmt.pf fmt "Sign"
  | ArgMax _ -> Fmt.pf fmt "ArgMax"
  | Pow _ -> Fmt.pf fmt "Pow"
  | QuantizeLinear _ -> Fmt.pf fmt "QuantizeLinear"
  | DequantizeLinear _ -> Fmt.pf fmt "DequantizeLinear"

let show_descr t = Fmt.str "%a" pp_descr t
let compare { id = id1; _ } { id = id2; _ } = Int.compare id1 id2
let equal { id = id1; _ } { id = id2; _ } = Int.equal id1 id2
let hash { id; _ } = id
let sexp_of_t node = Base.Int.sexp_of_t node.id
let pp fmt n = Fmt.pf fmt "@[%i: %a@]" n.id pp_descr n.descr
let show n = Fmt.str "%a" pp n

include Base.Comparator.Make (struct
  type nonrec t = t

  let compare = compare
  let sexp_of_t = sexp_of_t
end)

let rec compute_shape n = n.shape

and compute_shape_descr descr =
  match descr with
  | Add { input1; _ }
  | Div { input1; _ }
  | Mul { input1; _ }
  | Sub { input1; _ }
  | Pow { input1; _ } ->
    compute_shape input1
  | Flatten { input; axis } ->
    (* (d_0 X d_1 … d_(axis-1), d_axis X d_(axis+1) … X dn). *)
    let shape = compute_shape input in
    let d1 = ref 1 in
    let d2 = ref 1 in
    for i = 0 to axis - 1 do
      d1 := !d1 * Shape.get shape i
    done;
    for i = axis to Shape.rank shape - 1 do
      d2 := !d2 * Shape.get shape i
    done;
    Shape.of_list [ !d1; !d2 ]
  | Input { shape } -> shape
  | ReLu { input } -> compute_shape input
  | Transpose { input; perm = [] } ->
    compute_shape input |> Shape.to_list |> List.rev |> Shape.of_list
  | Transpose { input; perm } ->
    let shape = compute_shape input in
    let rank = Shape.rank shape in
    assert (Int.equal rank (List.length perm));
    let shape' = Array.create ~len:rank 0 in
    let shape = Shape.to_array shape in
    Base.List.iteri perm ~f:(fun i j -> Array.set shape' i (Array.get shape j));
    Shape.of_array shape'
  | Constant { data } -> Gentensor.shape data
  | Concat { inputs; axis } ->
    let shapes = List.map ~f:compute_shape inputs in
    let shape = List.hd_exn shapes in
    let axis = if axis < 0 then Shape.rank shape + axis else axis in
    let l = List.map ~f:(fun s -> Shape.get s axis) shapes in
    let i = List.reduce_exn ~f:( + ) l in
    Shape.set shape axis i
  | Gather { input; indices; axis } -> (
    let input_shape = compute_shape input in
    let indices_shape = compute_shape indices in
    let axis = if axis < 0 then Shape.rank input_shape + axis else axis in
    match List.split_n (Shape.to_list input_shape) axis with
    | _, [] ->
      Logging.user_error (fun m ->
        m "Axis is bigger than shape rank (%a)" pp_descr descr)
    | before, _ :: after ->
      Shape.of_list (before @ Shape.to_list indices_shape @ after))
  | Matmul { input1; input2 }
  | QLinearMatMul { inputA = input1; inputB = input2; _ } ->
    let pad_left = function
      | [] ->
        Logging.user_error (fun m ->
          m "Impossible to pad empty shape (%a)" pp_descr descr)
      | [ a ] -> ([ 1; a ], true)
      | x -> (x, false)
    in
    let rec remove_pad_left = function
      | [] ->
        Logging.user_error (fun m ->
          m "Impossible to remove pad empty shape (%a)" pp_descr descr)
      | [ k; a ] ->
        assert (k = 1);
        [ a ]
      | a :: l ->
        (* broadcasting can have been added *)
        a :: remove_pad_left l
    in
    let pad_right = function
      | [] ->
        Logging.user_error (fun m ->
          m "Impossible to pad empty shape (%a)" pp_descr descr)
      | [ a ] -> ([ a; 1 ], true)
      | x -> (x, false)
    in
    let rec remove_pad_right = function
      | [] ->
        Logging.user_error (fun m ->
          m "Impossible to remove pad empty shape (%a)" pp_descr descr)
      | [ a; k ] ->
        assert (k = 1);
        [ a ]
      | a :: l ->
        (* broadcasting can have been added *)
        a :: remove_pad_right l
    in
    let rec one_padding l i =
      if i <= 0 then l else one_padding (1 :: l) (i - 1)
    in
    (* Expected semantic is matrix multiplication C = AB. A (shape [n;m]); B
       (shape [m;p]); C (shape [n;p]) *)
    let check_matmul_size_ab ~a_sh ~(* shape A *) b_sh (* shape B *) =
      let adim2, pad_left_done = pad_left a_sh in
      let bdim2, pad_right_done = pad_right b_sh in
      let adim = one_padding adim2 (List.length bdim2 - List.length adim2) in
      let bdim = one_padding bdim2 (List.length adim2 - List.length bdim2) in
      let rec infer_csize acc ad bd =
        match (ad, bd) with
        | [ m; n ], [ nn; p ] ->
          if nn = n
          then List.rev_append acc [ m; p ]
          else
            Logging.user_error (fun m ->
              m "Size of matrices not adequate (%a)" pp_descr descr)
        | a :: la, b :: lb ->
          if a = b
          then infer_csize (a :: acc) la lb
          else if a = 1
          then infer_csize (b :: acc) la lb
          else if b = 1
          then infer_csize (a :: acc) la lb
          else
            Logging.user_error (fun m ->
              m "Checking size failed (%a): one discordance" pp_descr descr)
        | _, _ ->
          Logging.user_error (fun m ->
            m "Checking size failed (%a)" pp_descr descr)
      in
      let cdim = infer_csize [] adim bdim in
      (* When both padding are added the result is [1;1] which become [1] *)
      if pad_left_done
      then remove_pad_left cdim
      else if pad_right_done
      then remove_pad_right cdim
      else cdim
    in
    Shape.of_list
      (check_matmul_size_ab
         ~a_sh:(Shape.to_list (compute_shape input1))
         ~b_sh:(Shape.to_list (compute_shape input2)))
  | Reshape { input; shape; _ } ->
    let shape =
      match shape.descr with
      | Constant { data = Int64 a } ->
        List.map ~f:Int64.to_int_exn (Tensor.flatten a)
      | _ ->
        (* Some constant propagation could be useful in some cases eg. patch-1
           VNNcomp *)
        Logging.user_error (fun m ->
          m "Non-constant shape not supported (%a)" pp_descr descr)
    in
    List.iter shape ~f:(function
      | -1 ->
        Logging.user_error (fun m ->
          m "Shape value -1 not supported (%a)" pp_descr descr)
      | 0 ->
        Logging.user_error (fun m ->
          m "Shape value 0 not supported (%a)" pp_descr descr)
      | _ -> ());
    let out = Shape.of_list shape in
    if Shape.size out <> Shape.size input.shape
    then
      Logging.user_error (fun m ->
        m
          "Shape of input and shape given have not the same number of elements \
           (%a)"
          pp_descr descr);
    out
  | Gemm { inputA; inputB; transA; transB; _ }
  | QGemm { inputA; inputB; transA; transB; _ } ->
    let rank2 i =
      match Shape.to_array_unsafe i.shape with
      | [| k; n |] -> (k, n)
      | _ ->
        Logging.user_error (fun m ->
          m "Generalized matrix multiplications expect input shape of size 2")
    in
    let tr trans (k, n) = if trans = 1 then (n, k) else (k, n) in
    let a1, a2 = tr transA @@ rank2 inputA in
    let b1, b2 = tr transB @@ rank2 inputB in
    if not (Int.equal a2 b1)
    then
      Logging.user_error (fun m ->
        m "(M:%i,K:%i) times (K:%i,N:%i) results into (M:%i,N:%i)" a1 a2 b1 b2
          a1 b2);
    Shape.of_array [| a1; b2 |]
  | QuantizeLinear { x; axis; _ } | DequantizeLinear { x; axis; _ } ->
    ignore axis (* TODO: Something similar to ArgMax? *);
    compute_shape x
  | ArgMax { input; keepdims = true; _ } -> compute_shape input
  | ArgMax { input; axis; _ } ->
    let shape = compute_shape input in
    let rank = Shape.rank shape in
    if axis < -rank || axis >= rank
    then
      Logging.user_error (fun m ->
        m "Incorrect axis parameter (%a)" pp_descr descr)
    else Shape.remove_row shape @@ ((axis + rank) % rank)
  | Sign { input } -> compute_shape input
  | LogSoftmax | Squeeze _ | MaxPool | Conv | Identity _ | RW_Linearized_ReLu
  | ReduceSum _ | GatherND _ | RandomNormal _ | Abs _ | Log _ ->
    Logging.not_implemented_yet (fun m ->
      m "Shape computation (%a)" pp_descr descr)

let compute_ty n : ty =
  let same_type n1 n2 =
    if not (equal_ty n1.ty n2.ty)
    then
      Logging.user_error (fun m ->
        m "Expected same type argument, got %a and %a" pp_ty n1.ty pp_ty n2.ty)
  in
  match n with
  | Constant { data = Float _ } -> Float
  | Constant { data = UInt8 _ } -> UInt8
  | Constant { data = Int8 _ } -> Int8
  | Constant { data = Int32 _ } -> Int32
  | Constant { data = Int64 _ } -> Int64
  | Add { input1; input2; _ }
  | Div { input1; input2; _ }
  | Mul { input1; input2; _ }
  | Sub { input1; input2; _ }
  | Pow { input1; input2; _ } ->
    same_type input1 input2;
    input1.ty
  | Flatten { input; _ } | Sign { input; _ } -> input.ty
  | ArgMax _ -> Int64
  | QGemm { y_zero_point; _ } ->
    Option.value_map y_zero_point ~default:Float ~f:(fun y_zero_point ->
      y_zero_point.ty)
  | QLinearMatMul { y_zero_point; _ } -> y_zero_point.ty
  | QuantizeLinear { y_zero_point; _ } ->
    Option.value_map y_zero_point ~default:UInt8 ~f:(fun y_zero_point ->
      y_zero_point.ty)
  | DequantizeLinear _ -> Float
  | Matmul _ | Gemm _ | LogSoftmax | ReLu _ | Transpose _ | Squeeze _ | MaxPool
  | Conv | Reshape _ | Identity _ | Input _ | RW_Linearized_ReLu | Concat _
  | Gather _ | ReduceSum _ | GatherND _ | RandomNormal _ | Abs _ | Log _ ->
    Float (* TODO *)

let create =
  let c = ref (-1) in
  fun descr ->
    Int.incr c;
    { id = !c; descr; shape = compute_shape_descr descr; ty = compute_ty descr }

let constant_int_array a =
  create (Constant { data = Gentensor.of_int64_array a })

let reshape shape node =
  if Shape.equal node.shape shape
  then node
  else
    create
      (Reshape
         {
           input = node;
           shape =
             constant_int_array
               (Array.map ~f:Int64.of_int @@ Shape.to_array shape);
         })

let gather_int_as_matmul input i =
  let input1 = reshape (Shape.of_array [| 1; Shape.size input.shape |]) input in
  let selector = Array.create ~len:(Shape.size input1.shape) Float.zero in
  Array.set selector i Float.one;
  let selector =
    Gentensor.Float
      (Tensor.of_array1
         (Shape.of_array [| Array.length selector; 1 |])
         (Bigarray.Array1.of_array Float64 C_layout selector))
  in
  let input2 = create (Constant { data = selector }) in
  let result = create (Matmul { input1; input2 }) in
  reshape (Shape.of_array [| 1 |]) result

let gather_int ?(encode = true) input i =
  if encode
  then gather_int_as_matmul input i
  else
    let indices =
      create (Constant { data = Gentensor.create_1_int64 (Int64.of_int i) })
    in
    create (Gather { input; indices; axis = 0 })

let gather_ints_as_matmul input is =
  let input1 = reshape (Shape.of_array [| 1; Shape.size input.shape |]) input in
  let size = Shape.size input1.shape in
  let len = List.length is in
  let selector = Array.create ~len:(size * len) Float.zero in
  List.iteri is ~f:(fun j i -> Array.set selector ((j * size) + i) Float.one);
  let selector =
    Gentensor.Float
      (Tensor.of_array1
         (Shape.of_array [| size; len |])
         (Bigarray.Array1.of_array Float64 C_layout selector))
  in
  let input2 = create (Constant { data = selector }) in
  let result = create (Matmul { input1; input2 }) in
  reshape (Shape.of_array [| len |]) result

let gather_ints ?(encode = true) input is =
  if encode
  then gather_ints_as_matmul input is
  else
    let indices =
      create
        (Constant
           {
             data =
               Gentensor.of_int64_array
                 (Array.map ~f:Int64.of_int (List.to_array is));
           })
    in
    create (Gather { input; indices; axis = 0 })

let mul_float input f =
  let input1 = reshape (Shape.of_array [| 1; 1 |]) input in
  let f = Array.create ~len:1 f in
  let f =
    Gentensor.Float
      (Tensor.of_array1
         (Shape.of_array [| Array.length f; 1 |])
         (Bigarray.Array1.of_array Float64 C_layout f))
  in
  let input2 = create (Constant { data = f }) in
  let result = create (Matmul { input1; input2 }) in
  reshape (Shape.of_array [| 1 |]) result

let div_float ?(encode = true) input f =
  if encode
  then
    let f = Float.one /. f in
    mul_float input f
  else
    let input1 = reshape (Shape.of_array [| 1; 1 |]) input in
    let f = Array.create ~len:1 f in
    let f =
      Gentensor.Float
        (Tensor.of_array1
           (Shape.of_array [| Array.length f; 1 |])
           (Bigarray.Array1.of_array Float64 C_layout f))
    in
    let input2 = create (Constant { data = f }) in
    let result = create (Div { input1; input2 }) in
    reshape (Shape.of_array [| 1 |]) result

let ( + ) input1 input2 = create @@ Add { input1; input2 }
let ( * ) input1 input2 = create @@ Mul { input1; input2 }
let ( - ) input1 input2 = create @@ Sub { input1; input2 }

let concat_0 = function
  | [ n ] -> n
  | [] -> failwith "empty concat"
  | inputs -> create (Concat { inputs; axis = 0 })

let preds node =
  (* Predecessors must appear in the same order as in the ONNX. *)
  match node.descr with
  | Constant _ | Input _ -> []
  | Add { input1; input2 }
  | Sub { input1; input2 }
  | Mul { input1; input2 }
  | Div { input1; input2 }
  | Matmul { input1; input2 }
  | Pow { input1; input2 } ->
    [ input1; input2 ]
  | QLinearMatMul
      {
        inputA;
        inputA_scale;
        inputA_zero_point;
        inputB;
        inputB_scale;
        inputB_zero_point;
        y_scale;
        y_zero_point;
      } ->
    [
      inputA;
      inputA_scale;
      inputA_zero_point;
      inputB;
      inputB_scale;
      inputB_zero_point;
      y_scale;
      y_zero_point;
    ]
  | Gather { input; indices; axis = _ } -> [ input; indices ]
  | GatherND { data; indices; batch_dims = _ } -> [ data; indices ]
  | ReLu { input } | Abs { input } | Log { input } -> [ input ]
  | Concat { inputs; axis = _ } -> inputs
  | ReduceSum { input; axes = Some x; _ } -> [ input; x ]
  | ReduceSum { input; axes = None; _ } -> [ input ]
  | RandomNormal _ -> []
  | Transpose { input; _ } -> [ input ]
  | Flatten { input; _ } -> [ input ]
  | Identity { input } -> [ input ]
  | Gemm { inputA; inputB; inputC = Some x; _ } -> [ inputA; inputB; x ]
  | Gemm { inputA; inputB; inputC = None; _ } -> [ inputA; inputB ]
  | QGemm
      {
        inputA;
        inputA_scale;
        inputA_zero_point;
        inputB;
        inputB_scale;
        inputB_zero_point;
        inputC = Some inputC;
        y_scale = Some y_scale;
        y_zero_point = Some y_zero_point;
        _;
      } ->
    [
      inputA;
      inputA_scale;
      inputA_zero_point;
      inputB;
      inputB_scale;
      inputB_zero_point;
      inputC;
      y_scale;
      y_zero_point;
    ]
  | QGemm
      {
        inputA;
        inputA_scale;
        inputA_zero_point;
        inputB;
        inputB_scale;
        inputB_zero_point;
        inputC = None;
        y_scale = None;
        y_zero_point = None;
        _;
      } ->
    [
      inputA;
      inputA_scale;
      inputA_zero_point;
      inputB;
      inputB_scale;
      inputB_zero_point;
    ]
  | QGemm
      {
        inputA;
        inputA_scale;
        inputA_zero_point;
        inputB;
        inputB_scale;
        inputB_zero_point;
        inputC = Some inputC;
        y_scale = None;
        y_zero_point = None;
        _;
      } ->
    [
      inputA;
      inputA_scale;
      inputA_zero_point;
      inputB;
      inputB_scale;
      inputB_zero_point;
      inputC;
    ]
  | QGemm _ ->
    Logging.code_error ~src:Logging.src_nir (fun m ->
      m "Unexpected QGemm optional input values")
  | Squeeze { data; _ } -> [ data ]
  | Reshape { input; shape; _ } -> [ input; shape ]
  | LogSoftmax | MaxPool | Conv | RW_Linearized_ReLu -> []
  | Sign { input } -> [ input ]
  | ArgMax { input; _ } -> [ input ]
  | QuantizeLinear { x; y_scale; y_zero_point = Some y_zero_point; _ } ->
    [ x; y_scale; y_zero_point ]
  | QuantizeLinear { x; y_scale; y_zero_point = None; _ } -> [ x; y_scale ]
  | DequantizeLinear { x; x_scale; x_zero_point = Some x_zero_point; _ } ->
    [ x; x_scale; x_zero_point ]
  | DequantizeLinear { x; x_scale; x_zero_point = None; _ } -> [ x; x_scale ]

let map f n =
  match n.descr with
  | Constant _ | Input _ -> n
  | Add { input1; input2 } ->
    create (Add { input1 = f input1; input2 = f input2 })
  | Sub { input1; input2 } ->
    create (Sub { input1 = f input1; input2 = f input2 })
  | Mul { input1; input2 } ->
    create (Mul { input1 = f input1; input2 = f input2 })
  | Div { input1; input2 } ->
    create (Div { input1 = f input1; input2 = f input2 })
  | Matmul { input1; input2 } ->
    create (Matmul { input1 = f input1; input2 = f input2 })
  | QLinearMatMul t ->
    create (QLinearMatMul { t with inputA = f t.inputA; inputB = f t.inputB })
  | ReLu { input } -> create (ReLu { input = f input })
  | Abs { input } -> create (Abs { input = f input })
  | Log { input } -> create (Log { input = f input })
  | RandomNormal _ as descr -> create descr
  | ReduceSum { input; axes; keepdims; noop_with_empty_axes } ->
    create (ReduceSum { input = f input; axes; keepdims; noop_with_empty_axes })
  | Gather { input; indices; axis } ->
    create (Gather { input = f input; indices = f indices; axis })
  | GatherND { data; indices; batch_dims } ->
    create (GatherND { data = f data; indices = f indices; batch_dims })
  | Transpose t -> create (Transpose { t with input = f t.input })
  | Flatten t -> create (Flatten { t with input = f t.input })
  | Identity { input } -> create (Identity { input = f input })
  | Concat { inputs; axis } ->
    create (Concat { inputs = List.map ~f inputs; axis })
  | Gemm t ->
    create
      (Gemm
         {
           t with
           inputA = f t.inputA;
           inputB = f t.inputB;
           inputC = Option.map t.inputC ~f;
         })
  | QGemm t ->
    create
      (QGemm
         {
           t with
           inputA = f t.inputA;
           inputB = f t.inputB;
           inputC = Option.map t.inputC ~f;
         })
  | Squeeze t -> create (Squeeze { t with data = f t.data })
  | Reshape t -> create (Reshape { t with input = f t.input })
  | LogSoftmax | MaxPool | Conv | RW_Linearized_ReLu -> n (* todo *)
  | Sign { input } -> create (Sign { input = f input })
  | ArgMax t -> create (ArgMax { t with input = f t.input })
  | QuantizeLinear t -> create (QuantizeLinear { t with x = f t.x })
  | DequantizeLinear t -> create (DequantizeLinear { t with x = f t.x })
  | Pow { input1; input2 } ->
    create (Pow { input1 = f input1; input2 = f input2 })

let replace_input f node =
  let h = Base.Hashtbl.create (module Base.Int) in
  let rec aux n =
    Base.Hashtbl.find_or_add h n.id ~default:(fun () ->
      match n.descr with Input _ -> f () | _ -> map aux n)
  in
  aux node

(* Iter on the nodes accessible from [node] ([node] comprised) without
   repetition. *)
let map_rec f node =
  let h = Base.Hashtbl.create (module Base.Int) in
  let rec aux n =
    Base.Hashtbl.find_or_add h n.id ~default:(fun () -> f (map aux n))
  in
  aux node

let iter_rec f node =
  let h = Base.Hashtbl.create (module Base.Int) in
  let rec aux n =
    Base.Hashtbl.find_or_add h n.id ~default:(fun () ->
      List.iter ~f:aux (preds n);
      f n)
  in
  aux node

let sum_list ?(shp = Shape.of_array [| 1 |]) ns =
  match ns with
  | [] -> create @@ Constant { data = Gentensor.create_const_float shp 0.0 }
  | hd :: tl -> List.fold tl ~init:hd ~f:( + )

let partial_dot_product ?shp arr1 arr2 first last =
  let ioob str = failwith @@ "Index out of bound for arr" ^ str in
  if last > Array.length arr1
  then ioob "1"
  else if last > Array.length arr2
  then ioob "2"
  else if last > first
  then
    let rec aux index acc =
      if index = last
      then acc
      else
        let acc = acc + (arr1.(index) * arr2.(index)) in
        aux Int.(index + 1) acc
    in
    aux Int.(first + 1) (arr1.(first) * arr2.(first))
  else
    (* nothing to include, returns a tensor of 0s. *)
    let actual_shape =
      if Array.length arr1 <> 0
      then compute_shape arr1.(0)
      else if Array.length arr2 <> 0
      then compute_shape arr2.(0)
      else
        match shp with
        | Some s -> s
        | None -> failwith "Cannot determine shape of tensor"
    in
    create @@ Constant { data = Gentensor.create_const_float actual_shape 0.0 }

let transpose perm id =
  match perm with
  | [] -> id |> List.of_array |> List.rev |> List.to_array
  | _ ->
    let id' = Array.create ~len:(Array.length id) (-1) in
    List.iteri ~f:(fun i permi -> id'.(i) <- id.(permi)) perm;
    id'

let untranspose perm id =
  match perm with
  | [] -> id |> List.of_array |> List.rev |> List.to_array
  | _ ->
    let id' = Array.create ~len:(Array.length id) (-1) in
    List.iteri ~f:(fun i permi -> id'.(permi) <- id.(i)) perm;
    id'

let flatten shp axis id =
  Int.(
    let r = Shape.rank shp in
    let axis = (axis + r) % r in
    let result = [| 0; 0 |] in
    for i = 0 to axis - 1 do
      result.(0) <- (result.(0) * Shape.get shp i) + id.(i)
    done;
    for i = axis to r - 1 do
      result.(1) <- (result.(1) * Shape.get shp i) + id.(i)
    done;
    result)

let unflatten shp axis id =
  let r = Shape.rank shp in
  let axis = Int.(axis + r) % r in
  let result = Array.create ~len:r (-1) in
  let rec aux shift current_index current_value =
    if current_index >= 0
    then (
      let dim = Shape.get shp Int.(shift + current_index) in
      let modu = current_value % dim in
      let rest = current_value / dim in
      result.(Int.(shift + current_index)) <- modu;
      aux shift Int.(current_index - 1) rest)
  in
  aux 0 Int.(axis - 1) id.(0);
  aux axis Int.(r - 1 - axis) id.(1);
  result

let _replace_gather n =
  match n.descr with
  | Gather { input; indices; axis } ->
    let index =
      match indices.descr with
      | Constant { data } -> (
        match data with
        | Float _ | Int8 _ | UInt8 _ | Int32 _ ->
          assert false (* indices need to be int64s *)
        | Int64 data -> (
          match Tensor.flatten data with
          | [ data ] -> Option.value ~default:0 (Int64.to_int data)
          | _ -> assert false (* only one index *)))
      | _ -> assert false (* indices need to be constant *)
    in
    let rank = Shape.rank n.shape in
    (* axis = - 2*)
    let axis = Int.((rank + axis) % rank) in
    (* normalised axis *)
    let current_shape = Shape.to_array input.shape in
    let new_shape =
      Array.init
        Int.(rank - 1)
        ~f:(fun i -> Int64.of_int @@ if i < axis then i else Int.(i + 1))
    in
    if Int.(axis = rank - 2)
    then
      let nb_rows = current_shape.(Int.(rank - 2)) in
      let mat =
        Array.create ~len:1
          (Array.init nb_rows ~f:(fun i -> if i = index then 0.0 else 1.0))
      in
      let mat = create @@ Constant { data = Gentensor.of_float_matrix mat } in
      let n = create @@ Matmul { input1 = mat; input2 = n } in
      create
      @@ Reshape
           {
             input = n;
             shape =
               create @@ Constant { data = Gentensor.of_int64_array new_shape };
           }
    else if Int.((rank + axis) % rank = rank - 1)
    then (
      Stdlib.Printf.printf "\n";
      assert false)
    else assert false (* can only deal with axis = rank-1 or rank-2 atm *)
  | _ -> n

let encode_qgemm = function
  | QGemm
      {
        inputA;
        inputA_scale;
        inputA_zero_point;
        inputB;
        inputB_scale;
        inputB_zero_point;
        inputC;
        y_scale;
        y_zero_point;
        alpha;
        transA;
        transB;
      } ->
    let inputA =
      create
        (DequantizeLinear
           {
             x = inputA;
             x_scale = inputA_scale;
             x_zero_point = Some inputA_zero_point;
             axis = 1 (* Default value. *);
           })
    in
    let inputB =
      create
        (DequantizeLinear
           {
             x = inputB;
             x_scale = inputB_scale;
             x_zero_point = Some inputB_zero_point;
             axis = 1 (* Default value. *);
           })
    in
    let inputC =
      Option.map inputC ~f:(fun inputC ->
        create
          (DequantizeLinear
             {
               x = inputC;
               x_scale = inputA_scale * inputB_scale;
               x_zero_point = None;
               axis = 1 (* Default value. *);
             }))
    in
    let mmABC =
      create
        (Gemm
           {
             inputA;
             inputB;
             inputC;
             alpha;
             beta = 1.0 (* Default value. TODO? *) *. alpha;
             transA;
             transB;
           })
    in
    let y_scale =
      (* The QGemm specification requires [y_scale] to be a scalar. If supplied,
         the result is quantized, otherwise the result is full precision, ie
         there is no scaling, which is achieved here by a scaling of 1. *)
      Option.value y_scale
        ~default:(create (Constant { data = Gentensor.create_1_float 1.0 }))
    in
    QuantizeLinear
      { x = mmABC; y_scale; y_zero_point; axis = 1 (* Default value. *) }
  | descr ->
    Logging.user_error (fun m ->
      m "Expecting QGemm, got %a instead" pp_descr descr)