Source file torch_bindings.ml

open Ctypes

module C (F : Cstubs.FOREIGN) = struct
  open F

  let manual_seed = foreign "at_manual_seed" (int64_t @-> returning void)
  let free = foreign "free" (ptr void @-> returning void)
  let get_num_threads = foreign "at_get_num_threads" (void @-> returning int)
  let set_num_threads = foreign "at_set_num_threads" (int @-> returning void)

  module Tensor = struct
    type t = unit ptr

    let t : t typ = ptr void
    let new_tensor = foreign "at_new_tensor" (void @-> returning t)

    let tensor_of_data =
      foreign
        "at_tensor_of_data"
        (ptr void
         (* data *)
         @-> ptr int64_t
         (* dims *)
         @-> int
         (* ndims *)
         @-> int
         (* element size in bytes *)
         @-> int
         (* kind *)
         @-> returning t)
    ;;

    let copy_data =
      foreign
        "at_copy_data"
        (t
         (* tensor *)
         @-> ptr void
         (* data *)
         @-> int64_t
         (* numel *)
         @-> int
         (* element size in bytes *)
         @-> returning void)
    ;;

    let copy_ = foreign "at_copy_" (t (* dst *) @-> t (* src *) @-> returning void)
    let set_data = foreign "at_set_data" (t (* dst *) @-> t (* src *) @-> returning void)

    let float_vec =
      foreign
        "at_float_vec"
        (ptr double (* values *)
         @-> int (* num values *)
         @-> int (* kind *)
         @-> returning t)
    ;;

    let int_vec =
      foreign
        "at_int_vec"
        (ptr int64_t
         (* values *)
         @-> int
         (* num values *)
         @-> int
         (* kind *)
         @-> returning t)
    ;;

    let device = foreign "at_device" (t @-> returning int)
    let defined = foreign "at_defined" (t @-> returning bool)
    let num_dims = foreign "at_dim" (t @-> returning int)
    let shape = foreign "at_shape" (t @-> ptr int (* dims *) @-> returning void)
    let scalar_type = foreign "at_scalar_type" (t @-> returning int)
    let backward = foreign "at_backward" (t @-> int @-> int @-> returning void)
    let requires_grad = foreign "at_requires_grad" (t @-> returning int)
    let grad_set_enabled = foreign "at_grad_set_enabled" (int @-> returning int)
    let get = foreign "at_get" (t @-> int @-> returning t)

    let double_value =
      foreign "at_double_value_at_indexes" (t @-> ptr int @-> int @-> returning float)
    ;;

    let int64_value =
      foreign "at_int64_value_at_indexes" (t @-> ptr int @-> int @-> returning int64_t)
    ;;

    let double_value_set =
      foreign
        "at_set_double_value_at_indexes"
        (t @-> ptr int @-> int @-> float @-> returning void)
    ;;

    let int64_value_set =
      foreign
        "at_set_int64_value_at_indexes"
        (t @-> ptr int @-> int @-> int64_t @-> returning void)
    ;;

    let fill_double = foreign "at_fill_double" (t @-> float @-> returning void)
    let fill_int64 = foreign "at_fill_int64" (t @-> int64_t @-> returning void)
    let print = foreign "at_print" (t @-> returning void)
    let to_string = foreign "at_to_string" (t @-> int @-> returning string)
    let free = foreign "at_free" (t @-> returning void)

    let run_backward =
      foreign
        "at_run_backward"
        (ptr t @-> int @-> ptr t @-> int @-> ptr t @-> int @-> int @-> returning void)
    ;;
  end

  module Scalar = struct
    type t = unit ptr

    let t : t typ = ptr void
    let int = foreign "ats_int" (int64_t @-> returning t)
    let float = foreign "ats_float" (float @-> returning t)
    let free = foreign "ats_free" (t @-> returning void)
  end

  module Serialize = struct
    let t = Tensor.t
    let save = foreign "at_save" (t @-> string @-> returning void)
    let load = foreign "at_load" (string @-> returning t)

    let save_multi =
      foreign
        "at_save_multi"
        (ptr t @-> ptr (ptr char) @-> int @-> string @-> returning void)
    ;;

    let load_multi =
      foreign
        "at_load_multi"
        (ptr t @-> ptr (ptr char) @-> int @-> string @-> returning void)
    ;;

    let load_multi_ =
      foreign
        "at_load_multi_"
        (ptr t @-> ptr (ptr char) @-> int @-> string @-> returning void)
    ;;

    let load_callback =
      foreign
        "at_load_callback"
        (string
         @-> static_funptr Ctypes.(string @-> t @-> returning void)
         @-> returning void)
    ;;
  end

  module Optimizer = struct
    type t = unit ptr

    let t : t typ = ptr void

    let adam =
      foreign "ato_adam" (float @-> float @-> float @-> float @-> float @-> returning t)
    ;;

    let rmsprop =
      foreign
        "ato_rmsprop"
        (float
         (* learning rate *)
         @-> float
         (* alpha *)
         @-> float
         (* eps *)
         @-> float
         (* weight decay *)
         @-> float
         (* momentum *)
         @-> int
         (* centered *)
         @-> returning t)
    ;;

    let sgd =
      foreign
        "ato_sgd"
        (float
         (* learning rate *)
         @-> float
         (* momentum *)
         @-> float
         (* dampening *)
         @-> float
         (* weight decay *)
         @-> bool
         (* nesterov *)
         @-> returning t)
    ;;

    let add_parameters =
      foreign "ato_add_parameters" (t @-> ptr Tensor.t @-> int @-> returning void)
    ;;

    let set_learning_rate =
      foreign "ato_set_learning_rate" (t @-> float @-> returning void)
    ;;

    let set_momentum = foreign "ato_set_momentum" (t @-> float @-> returning void)
    let zero_grad = foreign "ato_zero_grad" (t @-> returning void)
    let step = foreign "ato_step" (t @-> returning void)
    let free = foreign "ato_free" (t @-> returning void)
  end

  module Cuda = struct
    let device_count = foreign "atc_cuda_device_count" (void @-> returning int)
    let is_available = foreign "atc_cuda_is_available" (void @-> returning int)
    let cudnn_is_available = foreign "atc_cudnn_is_available" (void @-> returning int)
    let set_benchmark_cudnn = foreign "atc_set_benchmark_cudnn" (int @-> returning void)
  end

  module Ivalue = struct
    type t = unit ptr

    let t : t typ = ptr void
    let to_int64 = foreign "ati_to_int" (t @-> returning int64_t)
    let to_bool = foreign "ati_to_bool" (t @-> returning int)
    let to_double = foreign "ati_to_double" (t @-> returning double)
    let to_tensor = foreign "ati_to_tensor" (t @-> returning Tensor.t)
    let tuple_length = foreign "ati_tuple_length" (t @-> returning int)
    let list_length = foreign "ati_list_length" (t @-> returning int)
    let to_tuple = foreign "ati_to_tuple" (t @-> ptr t @-> int @-> returning void)

    let to_generic_list =
      foreign "ati_to_generic_list" (t @-> ptr t @-> int @-> returning void)
    ;;

    let to_string = foreign "ati_to_string" (t @-> returning string)
    let none = foreign "ati_none" (void @-> returning t)
    let bool = foreign "ati_bool" (int @-> returning t)
    let tensor = foreign "ati_tensor" (Tensor.t @-> returning t)
    let int64 = foreign "ati_int" (int64_t @-> returning t)
    let double = foreign "ati_double" (float @-> returning t)
    let tuple = foreign "ati_tuple" (ptr t @-> int @-> returning t)
    let string = foreign "ati_string" (string @-> returning t)
    let tag = foreign "ati_tag" (t @-> returning int)
    let free = foreign "ati_free" (t @-> returning void)
  end

  module Module = struct
    type t = unit ptr

    let t : t typ = ptr void
    let load = foreign "atm_load" (string @-> returning t)
    let load_str = foreign "atm_load_str" (string @-> int @-> returning t)
    let forward = foreign "atm_forward" (t @-> ptr Tensor.t @-> int @-> returning Tensor.t)

    let forward_ =
      foreign "atm_forward_" (t @-> ptr Ivalue.t @-> int @-> returning Ivalue.t)
    ;;

    let named_buffers = foreign "atm_named_buffers" (t @-> returning Ivalue.t)
    let free = foreign "atm_free" (t @-> returning void)
  end

  module TensorG = Torch_bindings_generated.C (F)
end