arrayjit 0.3.3 · OCaml Package

module Lazy = Utils.Lazy

module Nd = Ndarray

module Debug_runtime = Utils.Debug_runtime

type work =

| Work of (module Minidebug_runtime.Debug_runtime) -> Base.unit -> Base.unit

val sexp_of_work : work -> Sexplib0.Sexp.t

val run : (module Minidebug_runtime.Debug_runtime) -> work -> 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 or merge call.
*)

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 jitted 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;

}

val sexp_of_t : t -> Sexplib0.Sexp.t

val name : t -> Base.String.t

val label : t -> Base.String.t

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

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

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.t

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

val hash : t -> Base.Hash.hash_value

val hash_fold_t : Base.Hash.state -> t -> Base.Hash.state

val hash_t : t -> Base.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 ident_label : t -> Base.String.t option

val debug_name : 
  id:Base.Int.t ->
  label:Base.String.t Base.List.t ->
  Base.String.t

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 | `Name_and_label | `Name_only ] ->
  t ->
  Base.String.t

val header : t -> string

module Registry : sig ... end

val registry : Registry.t

val create : 
  Ops.prec ->
  id:Base.int ->
  label:Base.string Base.list ->
  dims:int Base.Array.t Lazy.t ->
  Ops.init_op ->
  t

val print_accessible_headers : unit -> Base.unit

val log_accessible_headers : unit -> unit

package arrayjit