package torch

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

Source file tensor.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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
include Base
include Torch_core.Wrapper.Tensor

let set_float2 t i j value = float_set t [ i; j ] value
let set_float1 t i value = float_set t [ i ] value
let set_int2 t i j value = int_set t [ i; j ] value
let set_int1 t i value = int_set t [ i ] value
let get_float2 t i j = float_get t [ i; j ]
let get_float1 t i = float_get t [ i ]
let get_int2 t i j = int_get t [ i; j ]
let get_int1 t i = int_get t [ i ]
let ( .%{} ) = int_get
let ( .%{}<- ) = int_set
let ( .%.{} ) = float_get
let ( .%.{}<- ) = float_set
let ( .%[] ) = get_int1
let ( .%[]<- ) = set_int1
let ( .%.[] ) = get_float1
let ( .%.[]<- ) = set_float1

let no_grad_ t ~f =
  if requires_grad t
  then (
    let t = set_requires_grad t ~r:false in
    Exn.protect
      ~f:(fun () -> f t)
      ~finally:(fun () -> ignore (set_requires_grad t ~r:true : t)))
  else f t

let no_grad f =
  let prev = grad_set_enabled false in
  Exn.protect ~f ~finally:(fun () -> ignore (grad_set_enabled prev : bool))

let zero_grad t =
  let grad = grad t in
  if defined grad
  then (
    ignore (detach_ grad : t);
    ignore (zero_ grad : t))

type create =
  ?requires_grad:bool
  -> ?kind:Torch_core.Kind.packed
  -> ?device:Device.t
  -> ?scale:float
  -> int list
  -> t

let type_ = kind
let to_type t ~type_ = totype t ~scalar_type:type_
let to_kind t ~kind = totype t ~scalar_type:kind

let to_device ?device t =
  match device with
  | None -> t
  | Some device -> to_ t ~device

let float_vec ?kind ?device dims = float_vec ?kind dims |> to_device ?device

let gen
    ~f
    ?(requires_grad = false)
    ?(kind = Torch_core.Kind.(T Float))
    ?(device = Device.Cpu)
    ?scale
    size
  =
  let t = f ~size ~options:(kind, device) in
  let t =
    Option.value_map
      scale
      ~f:(fun scale -> mul t (float_vec [ scale ] ~device))
      ~default:t
  in
  if requires_grad then set_requires_grad t ~r:true else t

let zeros = gen ~f:zeros
let ones = gen ~f:ones
let rand = gen ~f:rand
let randn = gen ~f:randn
let f v = float_vec [ v ] |> reshape ~shape:[]
let mm = matmul
let ( + ) = add
let ( - ) = sub
let ( * ) = mul
let ( / ) = div
let ( ~- ) = neg
let ( -= ) t other = ignore (sub_ t other : t)
let ( += ) t other = ignore (add_ t other : t)
let ( /= ) t other = ignore (div_ t other : t)
let ( *= ) t other = ignore (mul_ t other : t)
let ( = ) = eq1
let pair_to_list (p1, p2) = [ p1; p2 ]

let conv2d ?(padding = 0, 0) ?(dilation = 1, 1) ?(groups = 1) input weight bias ~stride =
  conv2d
    input
    ~weight
    ~bias
    ~stride:(pair_to_list stride)
    ~padding:(pair_to_list padding)
    ~dilation:(pair_to_list dilation)
    ~groups

let conv_transpose2d
    ?(output_padding = 0, 0)
    ?(padding = 0, 0)
    ?(dilation = 1, 1)
    ?(groups = 1)
    input
    weight
    bias
    ~stride
  =
  conv_transpose2d
    input
    ~weight
    ~bias
    ~stride:(pair_to_list stride)
    ~padding:(pair_to_list padding)
    ~output_padding:(pair_to_list output_padding)
    ~groups
    ~dilation:(pair_to_list dilation)

let max_pool2d
    ?(padding = 0, 0)
    ?(dilation = 1, 1)
    ?(ceil_mode = false)
    ?stride
    self
    ~ksize
  =
  max_pool2d
    self
    ~kernel_size:(pair_to_list ksize)
    ~stride:(Option.value stride ~default:ksize |> pair_to_list)
    ~padding:(pair_to_list padding)
    ~dilation:(pair_to_list dilation)
    ~ceil_mode

let avg_pool2d
    ?(padding = 0, 0)
    ?(count_include_pad = false)
    ?(ceil_mode = false)
    ?stride
    ?divisor_override
    self
    ~ksize
  =
  let k1, k2 = ksize in
  let divisor_override = Option.value divisor_override ~default:Int.(k1 * k2) in
  avg_pool2d
    self
    ~kernel_size:(pair_to_list ksize)
    ~stride:(Option.value stride ~default:ksize |> pair_to_list)
    ~padding:(pair_to_list padding)
    ~ceil_mode
    ~count_include_pad
    ~divisor_override

let const_batch_norm ?(momentum = 0.1) ?(eps = 1e-5) input =
  batch_norm
    input
    ~weight:None
    ~bias:None
    ~running_mean:None
    ~running_var:None
    ~training:true
    ~momentum
    ~eps
    ~cudnn_enabled:false

let to_bigarray t ~kind =
  let bigarray = Bigarray.Genarray.create kind C_layout (shape t |> Array.of_list) in
  copy_to_bigarray (to_device t ~device:Cpu) bigarray;
  bigarray

let nll_loss ?(reduction = Torch_core.Reduction.Elementwise_mean) xs ~targets =
  nll_loss xs ~target:targets ~weight:None ~reduction ~ignore_index:(-100)

let cross_entropy_for_logits ?reduction logits ~targets =
  nll_loss ?reduction (log_softmax logits ~dim:(-1) ~dtype:(T Float)) ~targets

let dropout t ~p ~is_training = dropout t ~p ~train:is_training

let bce_loss ?(reduction = Torch_core.Reduction.Elementwise_mean) t ~targets =
  binary_cross_entropy t ~target:targets ~weight:None ~reduction

let mse_loss ?(reduction = Torch_core.Reduction.Elementwise_mean) t1 t2 =
  mse_loss t1 ~target:t2 ~reduction

let huber_loss ?(reduction = Torch_core.Reduction.Elementwise_mean) t1 t2 =
  let d = abs (t1 - t2) in
  let half = f 0.5 in
  let err = where1 ~condition:(le d (Scalar.float 1.)) (half * d * d) (d - half) in
  match reduction with
  | None -> err
  | Elementwise_mean -> mean err
  | Sum -> sum err

let bce_loss_with_logits ?(reduction = Torch_core.Reduction.Elementwise_mean) t ~targets =
  let max_val = clamp_min_ (-t) ~min:(Scalar.float 0.) in
  let one_minus_targets = ones_like targets - targets in
  let bce =
    add_
      (add_ (mul_ one_minus_targets t) max_val)
      (add_ (exp_ (-max_val)) (exp_ (-t - max_val)) |> log_)
  in
  match reduction with
  | None -> bce
  | Elementwise_mean -> mean bce
  | Sum -> sum bce

let pp formatter t =
  let shape = shape t in
  let element_count = List.fold shape ~init:1 ~f:Int.( * ) in
  if element_count < 1_000
  then (
    Caml.Format.pp_print_newline formatter ();
    Caml.Format.pp_print_string formatter (to_string t ~line_size:96);
    Caml.Format.pp_print_newline formatter ())
  else
    List.map shape ~f:Int.to_string
    |> String.concat ~sep:", "
    |> Printf.sprintf "Tensor<%s>"
    |> Caml.Format.pp_print_string formatter

let copy t =
  let t_ = zeros (shape t) ~kind:(kind t) in
  copy_ t_ ~src:t;
  t_

let shape_str t = List.map (shape t) ~f:Int.to_string |> String.concat ~sep:", "
let print_shape ?(name = "") t = Stdio.printf "%s<%s>\n%!" name (shape_str t)

let bigarray_to_array1 bigarray ~f =
  try
    let bigarray = Bigarray.array1_of_genarray bigarray in
    Array.init (Bigarray.Array1.dim bigarray) ~f:(fun i -> f bigarray.{i}) |> Option.some
  with
  | Invalid_argument _ -> None

let bigarray_to_array2 bigarray ~f =
  try
    let bigarray = Bigarray.array2_of_genarray bigarray in
    Array.init (Bigarray.Array2.dim1 bigarray) ~f:(fun i ->
        Array.init (Bigarray.Array2.dim2 bigarray) ~f:(fun j -> f bigarray.{i, j}))
    |> Option.some
  with
  | Invalid_argument _ -> None

let bigarray_to_array3 bigarray ~f =
  try
    let bigarray = Bigarray.array3_of_genarray bigarray in
    Array.init (Bigarray.Array3.dim1 bigarray) ~f:(fun i ->
        Array.init (Bigarray.Array3.dim2 bigarray) ~f:(fun j ->
            Array.init (Bigarray.Array3.dim3 bigarray) ~f:(fun k -> f bigarray.{i, j, k})))
    |> Option.some
  with
  | Invalid_argument _ -> None

let to_float1 t =
  match kind t with
  | T Float -> to_bigarray t ~kind:Bigarray.float32 |> bigarray_to_array1 ~f:Fn.id
  | T Double -> to_bigarray t ~kind:Bigarray.float64 |> bigarray_to_array1 ~f:Fn.id
  | _ -> None

let to_float2 t =
  match kind t with
  | T Float -> to_bigarray t ~kind:Bigarray.float32 |> bigarray_to_array2 ~f:Fn.id
  | T Double -> to_bigarray t ~kind:Bigarray.float64 |> bigarray_to_array2 ~f:Fn.id
  | _ -> None

let to_float3 t =
  match kind t with
  | T Float -> to_bigarray t ~kind:Bigarray.float32 |> bigarray_to_array3 ~f:Fn.id
  | T Double -> to_bigarray t ~kind:Bigarray.float64 |> bigarray_to_array3 ~f:Fn.id
  | _ -> None

let to_int1 t =
  match kind t with
  | T Int -> to_bigarray t ~kind:Bigarray.int32 |> bigarray_to_array1 ~f:Int32.to_int_exn
  | T Int64 ->
    to_bigarray t ~kind:Bigarray.int64 |> bigarray_to_array1 ~f:Int64.to_int_exn
  | _ -> None

let to_int2 t =
  match kind t with
  | T Int -> to_bigarray t ~kind:Bigarray.int32 |> bigarray_to_array2 ~f:Int32.to_int_exn
  | T Int64 ->
    to_bigarray t ~kind:Bigarray.int64 |> bigarray_to_array2 ~f:Int64.to_int_exn
  | _ -> None

let to_int3 t =
  match kind t with
  | T Int -> to_bigarray t ~kind:Bigarray.int32 |> bigarray_to_array3 ~f:Int32.to_int_exn
  | T Int64 ->
    to_bigarray t ~kind:Bigarray.int64 |> bigarray_to_array3 ~f:Int64.to_int_exn
  | _ -> None

let to_int1_exn t = Option.value_exn (to_int1 t)
let to_int2_exn t = Option.value_exn (to_int2 t)
let to_int3_exn t = Option.value_exn (to_int3 t)
let to_float1_exn t = Option.value_exn (to_float1 t)
let to_float2_exn t = Option.value_exn (to_float2 t)
let to_float3_exn t = Option.value_exn (to_float3 t)
let to_float0_exn = float_value

let to_float0 t =
  try float_value t |> Option.some with
  | _ -> None

let to_int0_exn = int_value

let to_int0 t =
  try int_value t |> Option.some with
  | _ -> None

let of_bigarray ?device ba = of_bigarray ba |> to_device ?device

let of_float0 ?device f =
  Bigarray.Array0.of_value Float32 C_layout f
  |> Bigarray.genarray_of_array0
  |> of_bigarray ?device

let of_float1 ?device f =
  Bigarray.Array1.of_array Float32 C_layout f
  |> Bigarray.genarray_of_array1
  |> of_bigarray ?device

let of_float2 ?device f =
  Bigarray.Array2.of_array Float32 C_layout f
  |> Bigarray.genarray_of_array2
  |> of_bigarray ?device

let of_float3 ?device f =
  Bigarray.Array3.of_array Float32 C_layout f
  |> Bigarray.genarray_of_array3
  |> of_bigarray ?device

let of_int0 ?device f =
  Bigarray.Array0.of_value Int C_layout f
  |> Bigarray.genarray_of_array0
  |> of_bigarray ?device

let of_int1 ?device f =
  Bigarray.Array1.of_array Int C_layout f
  |> Bigarray.genarray_of_array1
  |> of_bigarray ?device

let of_int2 ?device f =
  Bigarray.Array2.of_array Int C_layout f
  |> Bigarray.genarray_of_array2
  |> of_bigarray ?device

let of_int3 ?device f =
  Bigarray.Array3.of_array Int C_layout f
  |> Bigarray.genarray_of_array3
  |> of_bigarray ?device

let min_values = amin
let max_values = amax
let minimum t = reshape t ~shape:[ -1 ] |> min_values ~dim:[ 0 ] ~keepdim:false
let maximum t = reshape t ~shape:[ -1 ] |> max_values ~dim:[ 0 ] ~keepdim:false

let flatten t =
  let batch_size = shape t |> List.hd_exn in
  view t ~size:[ batch_size; -1 ]

let squeeze_last t = squeeze1 t ~dim:(-1)
let scale t f = mul1 t (Scalar.float f)
let eq_scalar = eq

let eq t1 t2 =
  if Torch_core.Kind.( <> ) (kind t1) (kind t2)
  then false
  else if Caml.( <> ) (shape t1) (shape t2)
  then false
  else eq1 t1 t2 |> all |> to_int0_exn |> ( <> ) 0

let to_list t =
  let size =
    match size t with
    | [] -> failwith "scalar tensor"
    | size :: _ -> size
  in
  List.init size ~f:(get t)
OCaml

Innovation. Community. Security.