Source file gentensor.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 t =
  | Float of (float, Bigarray.float64_elt) Tensor.t
  | Int8 of (int, Bigarray.int8_signed_elt) Tensor.t
  | UInt8 of (int, Bigarray.int8_unsigned_elt) Tensor.t
  | Int32 of (int32, Bigarray.int32_elt) Tensor.t
  | Int64 of (int64, Bigarray.int64_elt) Tensor.t

let create_1_float f = Float (Tensor.create_1_float f)
let create_1_int64 i = Int64 (Tensor.create_1_int64 i)
let create_const_float shape f = Float (Tensor.create_const_float shape f)

let of_int64_array ?shape t =
  let sh =
    match shape with
    | Some sh -> sh
    | None -> Shape.of_array [| Array.length t |]
  in
  let a = Bigarray.Array1.of_array Bigarray.Int64 Bigarray.c_layout t in
  Int64 (Tensor.of_array1 sh a)

let of_float_array ?shape t =
  let sh =
    match shape with
    | Some sh -> sh
    | None -> Shape.of_array [| Array.length t |]
  in
  let a = Bigarray.Array1.of_array Bigarray.Float64 Bigarray.c_layout t in
  Float (Tensor.of_array1 sh a)

let of_float_matrix ?(trans = false) m =
  let shape =
    if trans
    then [| Array.length m.(0); Array.length m |]
    else [| Array.length m; Array.length m.(0) |]
  in
  let tensor =
    Bigarray.Genarray.create Bigarray.Float64 Bigarray.c_layout shape
  in
  let position = [| 0; 0 |] in
  (* allocating this once and for all *)
  let first_index, second_index = if trans then (1, 0) else (0, 1) in
  for i = 0 to Array.length m - 1 do
    position.(first_index) <- i;
    let row = m.(i) in
    for j = 0 to Array.length row - 1 do
      position.(second_index) <- j;
      Bigarray.Genarray.set tensor position row.(j)
    done
  done;
  Float (Tensor.of_tensor tensor)

let of_reversed_float_matrix m =
  let tensor =
    Bigarray.Genarray.create Bigarray.Float64 Bigarray.c_layout
      [| Array.length m.(0); Array.length m |]
  in
  let position = [| 0; 0 |] in
  (* allocating this once and for all *)
  for i = 0 to Array.length m - 1 do
    position.(1) <- i;
    let row = m.(i) in
    for j = 0 to Array.length row - 1 do
      position.(0) <- j;
      Bigarray.Genarray.set tensor position row.(j)
    done
  done;
  Float (Tensor.of_tensor tensor)

let shape = function
  | Float f -> Tensor.shape f
  | Int8 i -> Tensor.shape i
  | UInt8 i -> Tensor.shape i
  | Int32 i -> Tensor.shape i
  | Int64 i -> Tensor.shape i