package rune

  1. Overview
  2. Docs
package rune
Legend:
Page
Library
Module
Module type
Parameter
Class
Class type
Source

Source file rune.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
include Tensor
include Tensor_with_debug

type ('a, 'b) t = ('a, 'b) Tensor.t
type float16_t = (float, float16_elt) t
type float32_t = (float, float32_elt) t
type float64_t = (float, float64_elt) t
type int8_t = (int, int8_elt) t
type uint8_t = (int, uint8_elt) t
type int16_t = (int, int16_elt) t
type uint16_t = (int, uint16_elt) t
type int32_t = (int32, int32_elt) t
type int64_t = (int64, int64_elt) t
type std_int_t = (int, int_elt) t
type std_nativeint_t = (nativeint, nativeint_elt) t
type complex32_t = (Complex.t, complex32_elt) t
type complex64_t = (Complex.t, complex64_elt) t

(* Re-export extended type aliases *)
type bfloat16_t = (float, Bigarray_ext.bfloat16_elt) t
type bool_t = (bool, Bigarray_ext.bool_elt) t
type int4_t = (int, Bigarray_ext.int4_signed_elt) t
type uint4_t = (int, Bigarray_ext.int4_unsigned_elt) t
type float8_e4m3_t = (float, Bigarray_ext.float8_e4m3_elt) t
type float8_e5m2_t = (float, Bigarray_ext.float8_e5m2_elt) t
type complex16_t = (Complex.t, Bigarray_ext.complex16_elt) t
type qint8_t = (int, Bigarray_ext.qint8_elt) t
type quint8_t = (int, Bigarray_ext.quint8_elt) t

(* Re-export extended dtype value constructors *)
let bfloat16 = Nx_core.Dtype.bfloat16
let bool = Nx_core.Dtype.bool
let int4 = Nx_core.Dtype.int4
let uint4 = Nx_core.Dtype.uint4
let float8_e4m3 = Nx_core.Dtype.float8_e4m3
let float8_e5m2 = Nx_core.Dtype.float8_e5m2
let complex16 = Nx_core.Dtype.complex16
let qint8 = Nx_core.Dtype.qint8
let quint8 = Nx_core.Dtype.quint8

(* ───── JIT ───── *)

type jit_device = [ `metal | `llvm ]

let is_jit_device_available = function
  | `llvm -> true
  | `metal -> (
      try
        let _ = Rune_jit_metal_or_missing.Device_info.get_default () in
        true
      with _ -> false)

let jit = Jit.jit

(* ───── Autodiff ───── *)

let grad = Autodiff.grad
let grads = Autodiff.grads
let value_and_grad = Autodiff.value_and_grad
let value_and_grads = Autodiff.value_and_grads
let jvp = Autodiff.jvp
let jvp_aux = Autodiff.jvp_aux
let jvps = Autodiff.jvps
let no_grad = Autodiff.no_grad
let detach = Autodiff.detach

(* ───── Gradient Checking ───── *)

module Finite_diff = Finite_diff
module Gradcheck = Gradcheck

type method_ = Finite_diff.method_

type gradient_check_result = Gradcheck.gradient_check_result = {
  max_abs_error : float;
  max_rel_error : float;
  mean_abs_error : float;
  mean_rel_error : float;
  failed_indices : (int array * float * float * float) list;
  passed : bool;
  num_checked : int;
  num_failed : int;
}

let finite_diff = Finite_diff.finite_diff
let finite_diff_jacobian = Finite_diff.finite_diff_jacobian
let check_gradient = Gradcheck.check_gradient
let check_gradients = Gradcheck.check_gradients

(* ───── Vmap ───── *)

type axis_spec = Vmap.axis_spec = Map of int | NoMap

type 'a in_axes_spec = 'a Vmap.in_axes_spec =
  | Single of axis_spec
  | Container of 'a

type 'a out_axes_spec = 'a Vmap.out_axes_spec =
  | OutSingle of int option
  | OutContainer of 'a

let vmap = Vmap.vmap
let vmaps = Vmap.vmaps

(* ───── RNG ───── *)

module Rng = Rng

(* ───── Debugging ───── *)

let debug = Debug.debug
let debug_with_context = Debug.with_context
let debug_push_context = Debug.push_context
let debug_pop_context = Debug.pop_context

(* ───── Nx Interop ───── *)

let of_nx nx_tensor = of_bigarray (Nx.to_bigarray nx_tensor)
let to_nx t = Nx.of_bigarray (to_bigarray t)