arrayjit 0.4.0 · OCaml Package

module Lazy = Utils.Lazy

module Nd = Ndarray

module Debug_runtime = Utils.Debug_runtime

val _get_local_debug_runtime : unit -> (module Minidebug_runtime.Debug_runtime)

type task =

| Task : {
1. context_lifetime : 'a;
2. description : Base.string;
3. work : Base.unit -> Base.unit;
} -> task

val sexp_of_task : task -> Sexplib0.Sexp.t

val run : task -> Base.unit

type memory_type =

| Constant
(*
The tensor node does not change after initialization.
*)
| Nonconstant
(*
One of: Changed_on_devices, Volatile.
*)
| Changed_on_devices
(*
The tensor node will only change on host via a to_host call.
*)
| Volatile
(*
The tensor node will only change on any device via a from_host call possibly followed by device_to_device.
*)

val memory_type_of_sexp : Sexplib0.Sexp.t -> memory_type

val sexp_of_memory_type : memory_type -> Sexplib0.Sexp.t

val compare_memory_type : memory_type -> memory_type -> Base.int

val equal_memory_type : memory_type -> memory_type -> Base.bool

type memory_mode =

| Effectively_constant
(*
Either Hosted Constant, or a subset of Virtual.
*)
| Virtual
(*
The tensor node's computations are inlined on a per-scalar basis.
*)
| Never_virtual
(*
One of: Local, On_device, Hosted.
*)
| Local
(*
The full tensor node is cached for the duration of a computation but not persisted across calls to compiled functions. It is not available for merging across devices.
*)
| Device_only
(*
One of: Local, On_device.
*)
| On_device
(*
The tensor node is stored on the devices that compute with it and persisted across function calls. It is available for merging across devices (for devices that support merging / P2P), but not (directly) for visualization or storing to disk.
*)
| Materialized
(*
One of: On_device, Hosted.
*)
| Hosted of memory_type
(*
The tensor node is stored in a globally addressable memory, in addition to on devices where it is computed with (or as part of one of them, if "hosting on device", or only on the host and not on devices, for some backends). It is available for all operations, and visible to OCaml programs as an Ndarray (the optional array of t).
*)

val memory_mode_of_sexp : Sexplib0.Sexp.t -> memory_mode

val sexp_of_memory_mode : memory_mode -> Sexplib0.Sexp.t

val compare_memory_mode : memory_mode -> memory_mode -> Base.int

val equal_memory_mode : memory_mode -> memory_mode -> Base.bool

type t = {

array : Nd.t Base.option Lazy.t;
prec : Ops.prec;
dims : Base.int Base.array Lazy.t;
id : Base.int;
label : Base.string Base.list;
(*
Display information. It is better if the last element of the list is the most narrow or alphanumeric, e.g. an identifier.
*)
mutable memory_mode : (memory_mode * Base.int) Base.option;
mutable backend_info : Base.Sexp.t;
mutable code_name : Base.string Base.option;

}

val sexp_of_t : t -> Sexplib0.Sexp.t

val compare : t -> t -> Base.int

val num_elems : t -> Base.int

val size_in_bytes : t -> Base__Int.t

val id : t -> Base.String.t

val label : t -> Base.String.t

val is_alphanum_ : Base.String.t -> bool

val get_debug_name : 
  ?code_name:Base.String.t ->
  id:Base__Int.t ->
  label:Base.String.t Base.List.t ->
  unit ->
  Base.String.t

val debug_name : t -> Base.String.t

val log_debug_info : from_log_level:int -> t -> unit

val default_to_most_local : t -> Base.int -> unit

val is_virtual_force : t -> Base.int -> bool

val is_hosted_force : ?specifically:memory_type -> t -> Base.int -> Base.bool

val is_materialized_force : t -> Base.int -> bool

val known_not_materialized : t -> bool

val known_constant : t -> bool

val known_non_virtual : t -> bool

val known_not_param : t -> bool

val mode_is_unspecified : t -> bool

val update_memory_mode : t -> memory_mode -> Base.int -> unit

include sig ... end

type comparator_witness

val comparator : (t, comparator_witness) Base__Comparator.comparator

val equal : t -> t -> Base.bool

val hash : t -> Base__Ppx_hash_lib.Std.Hash.hash_value

val hash_fold_t : 
  Base__.Ppx_hash_lib.Std.Hash.state ->
  t ->
  Base__.Ppx_hash_lib.Std.Hash.state

val hash_t : t -> Base__Ppx_hash_lib.Std.Hash.hash_value

val get_exn : t -> Nd.t

val has : t -> bool

val dims_to_string : ?with_axis_numbers:bool -> t -> Base.String.t

val no_grad_ident_label : t -> bool * Base.String.t option

val styled_ident : 
  repeating_nograd_idents:(Base.String.t, 'a) Base.Hashtbl.t ->
  repeating_grad_idents:(Base.String.t, 'a) Base.Hashtbl.t ->
  [< `Heuristic_ocannl of [< `Dot_grad | `Under_grad ]
  | `Name_and_label
  | `Name_only ] ->
  t ->
  Base.String.t

val update_code_name : t -> Base.string -> unit

val get_style : 
  ?arg_name:Base.String.t ->
  ?no_dots:bool ->
  unit ->
  [> `Heuristic_ocannl of [> `Dot_grad | `Under_grad ]
  | `Name_and_label
  | `Name_only ]

val header : t -> string

module Registry : sig ... end

val registry : Registry.t

val create : 
  Ops.prec ->
  id:Base__Int.t ->
  label:Base.String.t Base.List.t ->
  dims:Base__Int.t Base.Array.t Lazy.t ->
  Ops.init_op ->
  t

val find : id:Base.int -> Registry.data option

val print_accessible_headers : unit -> Base.unit

val log_accessible_headers : unit -> unit

package arrayjit