Something went wrong. Try again.
forked from raven-ml/raven
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471(*--------------------------------------------------------------------------- Copyright (c) 2026 The Raven authors. All rights reserved. SPDX-License-Identifier: ISC ---------------------------------------------------------------------------*)
(* Nx's elementwise functions on the host: each computes every dtype its kind takes, which its doc states, as the host's kernel called once on the operands broadcast by hand, whose kinds the kernel suite holds against nx_kinds.h; a dtype the kind does not take raises naming the function. *)
open Windtrapopen Nx_supportmodule A = Nx_arraymodule D = Nx_array.Dtypemodule L = Nx_array.Layoutmodule M = Nx_array.Movemodule P = Nx_kernel.Progmodule Support = Nx_kernels_support
type host = Nx.host
(* Arrays of any dtype *)
type case = Case : ('v, 's) A.t -> case
let pp_case ppf (Case a) = Format.fprintf ppf "%a %a" D.pp (A.dtype a) L.pp (A.layout a)
(* A contiguous array of [dt] and shape [s] over drawn bytes, a boolean's [0] or [1]. *)let drawn (type v s) (dt : (v, s) D.t) s : (v, s) A.t Gen.t = let open Gen in let n = Array.fold_left ( * ) 1 s in let+ data = string_of ~size:(constant (D.bytes dt n)) char in let data = match dt with | D.Bool -> String.map (fun c -> Char.chr (Char.code c land 1)) data | _ -> data in A.v dt (L.contiguous (Iarray.of_array s)) (Rig.Buffer.of_string (if data = "" then "\000" else data))
let shape = Gen.array ~size:(Gen.int_range 0 3) (Gen.frequency [ (1, Gen.constant 0); (6, Gen.int_range 1 4) ])
(* Every dtype, the complex ones, 2 of [D.all]'s, drawn a quarter of the time so that each run reaches the laws' complex cases. *)let dtype = let complex = List.filter (fun (D.Any dt) -> D.is D.Complex dt) D.all in Gen.with_pp (fun ppf (D.Any dt) -> D.pp ppf dt) (Gen.frequency [ (5, Gen.of_list D.all); (1, Gen.of_list complex) ])
(* [s] with some extents [1] and some leading axes dropped: a shape that broadcasts to [s]. *)let narrower s = let open Gen in let* drop = int_range 0 (Array.length s) in let+ ones = array ~size:(constant (Array.length s - drop)) bool in Array.mapi (fun i one -> if one then 1 else s.(i + drop)) ones
(* Operands of one dtype: the first of [s], each other of [s] or a shape that broadcasts to it, [n] in all. *)let operands n = let open Gen in let* (D.Any dt) = dtype in let* s = shape in let* first = drawn dt s in let+ rest = list ~size:(constant (n - 1)) (let* s' = frequency [ (1, constant s); (1, narrower s) ] in drawn dt s') in (s, Case first, List.map (fun a -> Case (A.expect dt (A.Any a))) rest)
let pp_operands ppf (s, c, cs) = Format.fprintf ppf "%a: %a" pp_ints s (Format.pp_print_list ~pp_sep:Format.pp_print_space pp_case) (c :: cs)
let unary_operands = Gen.with_pp pp_operands (operands 1)let binary_operands = Gen.with_pp pp_operands (operands 2)let ternary_operands = Gen.with_pp pp_operands (operands 3)
(* Bits *)
(* [a]'s bits in C order: one code per element of a sub-byte dtype, the bytes of any other. *)let bits_of (type v s) (a : (v, s) A.t) = let a = A.copy a in match D.bits (A.dtype a) with | 1 -> Array.map Bool.to_int (A.to_array (A.expect D.Bit (A.Any a))) | 4 -> A.to_array (Option.get (A.bitcast D.Uint4 a)) | _ -> A.to_array (Option.get (A.bitcast D.Uint8 a))
let value a : (_, _, host) Nx.t = Nx.Repr.of_array Nx.Host.v alet host_array x = Option.get (Nx.Repr.array (Nx.place Nx.Host.on x))let result_bits x = bits_of (host_array x)
(* [a] broadcast to [s], as a view. *)let stretched s a = if Iarray.to_array (L.shape (A.layout a)) = s then a else Option.get (A.move (M.Broadcast (Iarray.of_array s)) a)
(* nx.cpu, whose one-node maps are the references. *)let cpu = (module Nx_cpu : Nx_kernel.S)
let reference1 k dt s a = let dst = A.create Rig.host dt (Iarray.of_array s) in match Support.map1 cpu k ~dst (stretched s a) with | A.Done -> Some (bits_of dst) | _ | (exception Invalid_argument _) -> None
let reference2 k dt s a b = let dst = A.create Rig.host dt (Iarray.of_array s) in match Support.map2 cpu k ~dst (stretched s a) (stretched s b) with | A.Done -> Some (bits_of dst) | _ | (exception Invalid_argument _) -> None
let reference3 k s c a b = let dst = A.create Rig.host (A.dtype a) (Iarray.of_array s) in match Support.map3 cpu k ~dst (stretched s c) (stretched s a) (stretched s b) with | A.Done -> Some (bits_of dst) | _ | (exception Invalid_argument _) -> None
let names name f = raises_match (Exn.invalid_arg ~substring:("Nx." ^ name ^ ": ")) f
(* The law: where its kind takes the dtype, [f] computes, giving the reference's bits; elsewhere it raises naming itself. *)let agrees name ~accepted reference f = if accepted then begin cover "computed" true; match reference () with | Some want -> equal ~msg:name (array int) want (f ()) | None -> fail (name ^ ": the host's kernel declined a dtype its kind takes") end else begin cover "raised" true; names name (fun () -> ignore (f ())) end
(* Unary *)
type unary = { name : string; kind : P.unary; f : 'v 's. ('v, 's, host) Nx.t -> ('v, 's, host) Nx.t;}
let unaries = [ { name = "neg"; kind = Neg; f = Nx.neg }; { name = "recip"; kind = Recip; f = Nx.recip }; { name = "abs"; kind = Abs; f = Nx.abs }; { name = "sign"; kind = Sign; f = Nx.sign }; { name = "sqrt"; kind = Sqrt; f = Nx.sqrt }; { name = "exp"; kind = Exp; f = Nx.exp }; { name = "exp2"; kind = Exp2; f = Nx.exp2 }; { name = "expm1"; kind = Expm1; f = Nx.expm1 }; { name = "log"; kind = Log; f = Nx.log }; { name = "log2"; kind = Log2; f = Nx.log2 }; { name = "log1p"; kind = Log1p; f = Nx.log1p }; { name = "sin"; kind = Sin; f = Nx.sin }; { name = "cos"; kind = Cos; f = Nx.cos }; { name = "tan"; kind = Tan; f = Nx.tan }; { name = "asin"; kind = Asin; f = Nx.asin }; { name = "acos"; kind = Acos; f = Nx.acos }; { name = "atan"; kind = Atan; f = Nx.atan }; { name = "sinh"; kind = Sinh; f = Nx.sinh }; { name = "cosh"; kind = Cosh; f = Nx.cosh }; { name = "tanh"; kind = Tanh; f = Nx.tanh }; { name = "floor"; kind = Floor; f = Nx.floor }; { name = "ceil"; kind = Ceil; f = Nx.ceil }; { name = "round"; kind = Round; f = Nx.round }; { name = "trunc"; kind = Trunc; f = Nx.trunc }; ]
let law_unary (s, Case a, _) = let dt = A.dtype a in List.iter (fun u -> agrees u.name ~accepted:(P.accepts1 (Unary u.kind) dt dt) (fun () -> reference1 (Unary u.kind) dt s a) (fun () -> result_bits (u.f (value a)))) unaries
(* [erf] takes floats alone, by its type. *)let law_erf (s, Case a, _) = let dt = A.dtype a in match D.kind dt with | D.Float -> agrees "erf" ~accepted:(P.accepts1 (Unary Erf) dt dt) (fun () -> reference1 (Unary Erf) dt s a) (fun () -> result_bits (Nx.erf (value a))) | D.Signed | D.Unsigned | D.Boolean | D.Complex -> cover "computed" true
(* [asinh], [acosh] and [atanh] compute complex numbers by their kinds, and floats by a composition of other kinds that test_math holds. *)let law_inverse_hyperbolic (s, Case a, _) = let dt = A.dtype a in List.iter (fun (name, kind, f) -> match D.kind dt with | D.Float -> cover "computed" true | _ -> agrees name ~accepted:(P.accepts1 (Unary kind) dt dt) (fun () -> reference1 (Unary kind) dt s a) (fun () -> result_bits (f (value a)))) [ ("asinh", P.Asinh, fun x -> Nx.asinh x); ("acosh", P.Acosh, fun x -> Nx.acosh x); ("atanh", P.Atanh, fun x -> Nx.atanh x); ]
(* Binary *)
type binary = { name : string; kind : D.any -> P.binary; f : 'v 's. ('v, 's, host) Nx.t -> ('v, 's, host) Nx.t -> ('v, 's, host) Nx.t;}
let fixed k _ = k
let binaries = [ { name = "add"; kind = fixed P.Add; f = Nx.add }; { name = "sub"; kind = fixed P.Sub; f = Nx.sub }; { name = "mul"; kind = fixed P.Mul; f = Nx.mul }; { name = "div"; kind = (fun (D.Any dt) -> match D.kind dt with D.Signed | D.Unsigned -> P.Idiv | _ -> P.Fdiv); f = Nx.div; }; { name = "mod_"; kind = fixed P.Mod; f = Nx.mod_ }; { name = "pow"; kind = fixed P.Pow; f = Nx.pow }; { name = "atan2"; kind = fixed P.Atan2; f = Nx.atan2 }; { name = "maximum"; kind = fixed P.Maximum; f = Nx.maximum }; { name = "minimum"; kind = fixed P.Minimum; f = Nx.minimum }; { name = "bitwise_and"; kind = fixed P.And; f = Nx.bitwise_and }; { name = "bitwise_or"; kind = fixed P.Or; f = Nx.bitwise_or }; { name = "bitwise_xor"; kind = fixed P.Xor; f = Nx.bitwise_xor }; ]
let law_binary (s, Case a, rest) = let dt = A.dtype a in let b = match rest with [ Case b ] -> A.expect dt (A.Any b) | _ -> a in cover "broadcast" (Iarray.to_array (L.shape (A.layout b)) <> s); cover "complex" (D.is D.Complex dt); List.iter (fun (o : binary) -> let k = P.Binary (o.kind (D.Any dt)) in agrees o.name ~accepted:(P.accepts2 k dt) (fun () -> reference2 k dt s a b) (fun () -> result_bits (o.f (value a) (value b)))) binaries
(* Comparisons *)
type comparison = { name : string; kind : P.compare; swapped : bool; f : 'v 's. ('v, 's, host) Nx.t -> ('v, 's, host) Nx.t -> (bool, D.bool_elt, host) Nx.t;}
let comparisons = [ { name = "equal"; kind = Equal; swapped = false; f = Nx.equal }; { name = "not_equal"; kind = Not_equal; swapped = false; f = Nx.not_equal }; { name = "less"; kind = Less; swapped = false; f = Nx.less }; { name = "less_equal"; kind = Less_equal; swapped = false; f = Nx.less_equal; }; { name = "greater"; kind = Less; swapped = true; f = Nx.greater }; { name = "greater_equal"; kind = Less_equal; swapped = true; f = Nx.greater_equal; }; ]
let law_compare (s, Case a, rest) = let dt = A.dtype a in cover "complex" (D.is D.Complex dt); let b = match rest with [ Case b ] -> A.expect dt (A.Any b) | _ -> a in List.iter (fun (o : comparison) -> let k = P.Compare o.kind in let x, y = if o.swapped then (b, a) else (a, b) in agrees o.name ~accepted:(P.accepts2 k dt) (fun () -> reference2 k D.Bool s x y) (fun () -> result_bits (o.f (value a) (value b)))) comparisons
(* Three operands *)
let law_fma (s, Case a, rest) = let dt = A.dtype a in let b, c = match rest with | [ Case b; Case c ] -> (A.expect dt (A.Any b), A.expect dt (A.Any c)) | _ -> (a, a) in agrees "fma" ~accepted:(P.accepts3 Fma dt dt) (fun () -> reference3 Fma s a b c) (fun () -> result_bits (Nx.fma (value a) (value b) (value c)))
let law_where ((s, Case a, rest), flags) = let dt = A.dtype a in let b = match rest with Case b :: _ -> A.expect dt (A.Any b) | [] -> a in let n = Array.fold_left ( * ) 1 s in let c = A.of_array D.Bool (Iarray.of_array s) (Array.init n (fun i -> List.nth flags (i mod 8))) in agrees "where" ~accepted:true (fun () -> reference3 Where s c a b) (fun () -> result_bits (Nx.where (value c) (value a) (value b)))
(* Bitcast *)
let law_bitcast_same (_, Case a, _) = let x = value a in equal bool true (Nx.bitcast (A.dtype a) x == x)
(* atan2's range ends: [-π] at [y = -0.], [x < 0], and [π] rounded to the dtype, 3.25 in float8_e4m3fn. *)let test_atan2_ends () = let at dt y x = Nx.atan2 (Nx.scalar dt y) (Nx.scalar dt x) in let first x = (A.to_array (host_array x)).(0) in equal float_exact (-.Float.pi) (first (at D.Float64 (-0.) (-1.))); equal float_exact Float.pi (first (at D.Float64 0. (-1.))); equal float_exact 3.25 (first (at D.Float8_e4m3fn 0. (-1.)))
(* Composites *)
(* lerp's two forms at float64, as nx.mli states them. *)let lerp_reference a b w = let d = b -. a in if w < 0.5 then Float.fma w d a else Float.fma (w -. 1.) d b
let same_bits a b = (Float.is_nan a && Float.is_nan b) || Int64.equal (Int64.bits_of_float a) (Int64.bits_of_float b)
let weight = Gen.frequency [ (3, Gen.float_range (-2.) 3.); (2, Gen.of_list [ 0.; -0.; 1.; 0.5; -0.5; Float.pred 0.5; 2. ]); (1, Gen.any_float); ]
let ends = Gen.frequency [ (3, Gen.float_range (-1e3) 1e3); (1, Gen.of_list [ 0.; -0. ]); (1, Gen.any_float); ]
(* Equal as numbers: [-0.] and [+0.] alike, as nx.mli states the ends. *)let number = Testable.make ~pp:(fun ppf x -> Format.fprintf ppf "%h" x) ~equal:Float.equal
let law_lerp ((a, b), w) = cover "an end" (w = 0. || w = 1.); cover "the form from b" (w >= 0.5); cover "below -1/2" (w <= -0.5); let s = Nx.scalar Nx.float64 in let got = (Nx.to_array (Nx.lerp (s a) (s b) (s w))).(0) in let want = lerp_reference a b w in if not (same_bits want got) then failf "lerp %h %h %h is %h, not %h" a b w got want; if Float.is_finite (b -. a) && w = 0. then equal ~msg:"w = 0" number a got; if Float.is_finite (b -. a) && w = 1. then equal ~msg:"w = 1" number b got
(* A narrow float computes at float32 and rounds once. *)
type narrow = Narrow : (float, 's) D.t -> narrow
let narrows = Gen.of_list ~pp:(fun ppf (Narrow dt) -> D.pp ppf dt) D. [ Narrow Float16; Narrow Bfloat16; Narrow Float8_e4m3fn; Narrow Float8_e5m2; Narrow Float4_e2m1fn; ]
let law_lerp_narrow (Narrow dt, ((a, b), w)) = let s x = Nx.scalar dt x in let a = s a and b = s b and w = s w in let wide x = Nx.cast Nx.float32 x in let want = Nx.cast dt (Nx.lerp (wide a) (wide b) (wide w)) in (* Each element's bits, through the exact cast to float64, [None] for a NaN. *) let bits x = Array.map (fun v -> if Float.is_nan v then None else Some (Int64.bits_of_float v)) (Nx.to_array (Nx.cast Nx.float64 x)) in equal (array (option int64)) (bits want) (bits (Nx.lerp a b w))
(* The ends at signed zeros: equal as numbers. *)let test_lerp_zero_ends () = let s = Nx.scalar Nx.float64 in let lerp a b w = (Nx.to_array (Nx.lerp (s a) (s b) (s w))).(0) in List.iter (fun (a, b) -> equal ~msg:"w = 0" number a (lerp a b 0.); equal ~msg:"w = 1" number b (lerp a b 1.)) [ (-0., 1.); (1., -0.); (-0., -0.); (0., -0.); (-0., 0.) ]
let test_lerp_shapes () = let a = Nx.zeros Nx.float32 [| 2; 1 |] and b = Nx.ones Nx.float32 [| 3 |] in let y = Nx.lerp a b (Nx.scalar Nx.float32 0.25) in equal (array int) [| 2; 3 |] (Nx.shape y); equal (array float_exact) (Array.make 6 0.25) (Nx.to_array y); raises_match (Exn.invalid_arg ~substring:"Nx.lerp: ") (fun () -> ignore (Nx.lerp b (Nx.ones Nx.float32 [| 4 |]) b))
let elementwise = group "elementwise" [ prop "each unary function is its kind" unary_operands law_unary; prop "each binary function is its kind, broadcast" binary_operands law_binary; prop "each comparison is its kind, broadcast" binary_operands law_compare; prop "fma is its kind, broadcast" ternary_operands law_fma; prop "where is its kind, broadcast" (Gen.pair ternary_operands (Gen.list ~size:(Gen.constant 8) Gen.bool)) law_where; prop "erf is its kind" unary_operands law_erf; prop "asinh, acosh and atanh of complex numbers are their kinds" unary_operands law_inverse_hyperbolic; prop "a bitcast to its own dtype is the value itself" unary_operands law_bitcast_same; test "atan2 reaches its range's ends, rounded" test_atan2_ends; prop "lerp is one fma, from a below w = 1/2" (Gen.pair (Gen.pair ends ends) weight) law_lerp; prop "lerp of a narrow float is lerp at float32, rounded once" (Gen.pair narrows (Gen.pair (Gen.pair ends ends) weight)) law_lerp_narrow; test "lerp's ends at signed zeros" test_lerp_zero_ends; test "lerp broadcasts, and refuses shapes that do not" test_lerp_shapes; ]
let () = exit (run "nx elementwise" [ elementwise ])