Source file gentensor.ml

(**************************************************************************)
(*                                                                        *)
(*  This file is part of CAISAR.                                          *)
(*                                                                        *)
(*  Copyright (C) 2024                                                    *)
(*    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
  | 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 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 shape = function Float f -> Tensor.shape f | Int64 i -> Tensor.shape i