Fixes #387: normal distribution

lukstafi · lukstafi · commit c88f239a38bd · 2025-09-23T22:39:20.000+02:00
diff --git a/bin/compilation_speed.ml b/bin/compilation_speed.ml
@@ -3,7 +3,7 @@ open Ocannl
 module Tn = Ir.Tnode
 module IDX = Train.IDX
 module CDSL = Train.CDSL
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 
 (* FIXME: expose backend by name from Context *)
 
diff --git a/lib/nn_blocks.ml b/lib/nn_blocks.ml
@@ -20,6 +20,22 @@ open! Base
 open Ocannl_tensor.Operation.DSL_modules
 module Tn = Ir.Tnode
 
+let%op box_muller grad_spec init_f () =
+  let epsilon = [%oc Float.ldexp 1. (-24)] in
+  let u1 = init_f () in
+  let u2 = init_f () in
+  Ocannl_tensor.Operation.pointmul ~grad_spec
+    (sqrt (-2. *. log (u1 + (!.epsilon *. (1. - u1)))))
+    (cos (2. *. !.Float.pi *. u2))
+
+[%%extend_dsls
+let normal () = [%oc box_muller grad_spec uniform ()]
+let normal1 () = [%oc box_muller grad_spec uniform1 ()]
+let normal_at counter = [%oc box_muller grad_spec (fun () -> uniform_at counter) ()]
+let normal_at1 counter = [%oc box_muller grad_spec (fun () -> uniform_at1 counter) ()]]
+
+open DSL_modules
+
 let%op mlp_layer ~label ~hid_dim () x = relu (({ w } * x) + { b = 0.; o = [ hid_dim ] })
 
 (** Masks and scales by 1/keep_prob to maintain expected value. When [train_step = None], the
diff --git a/test/einsum/einsum_trivia.ml b/test/einsum/einsum_trivia.ml
@@ -3,7 +3,7 @@ open Ocannl
 module Tn = Ir.Tnode
 module IDX = Train.IDX
 module CDSL = Train.CDSL
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 
 module type Backend = Ir.Backend_intf.Backend
 
diff --git a/test/einsum/einsum_trivia_exec.ml b/test/einsum/einsum_trivia_exec.ml
@@ -3,7 +3,7 @@ open Ocannl
 module Tn = Ir.Tnode
 module IDX = Train.IDX
 module CDSL = Train.CDSL
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 
 module type Backend = Ir.Backend_intf.Backend
 
diff --git a/test/einsum/inline_permuted_view.ml b/test/einsum/inline_permuted_view.ml
@@ -2,7 +2,7 @@ open Base
 open Ocannl
 module IDX = Train.IDX
 module CDSL = Train.CDSL
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 
 module type Backend = Ir.Backend_intf.Backend
 
diff --git a/test/einsum/moons_demo_variant.ml b/test/einsum/moons_demo_variant.ml
@@ -1,7 +1,7 @@
 open Base
 open Ocannl
 module IDX = Train.IDX
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 module CDSL = Train.CDSL
 module Asgns = Ir.Assignments
 module Tn = Ir.Tnode
diff --git a/test/einsum/surjectivity.ml b/test/einsum/surjectivity.ml
@@ -2,7 +2,7 @@ open Base
 open Ocannl
 module IDX = Train.IDX
 module CDSL = Train.CDSL
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 
 module type Backend = Ir.Backend_intf.Backend
 
diff --git a/test/einsum/test_accumulation_semantics.ml b/test/einsum/test_accumulation_semantics.ml
@@ -2,7 +2,7 @@ open Base
 open Ocannl
 module IDX = Train.IDX
 module CDSL = Train.CDSL
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 
 module type Backend = Ir.Backend_intf.Backend
 
diff --git a/test/einsum/test_einsum_capture.ml b/test/einsum/test_einsum_capture.ml
@@ -2,7 +2,7 @@ open Base
 open Ocannl
 
 let capture_for_computation () =
-  let open Operation.DSL_modules in
+  let open Nn_blocks.DSL_modules in
   Tensor.unsafe_reinitialize ();
   let ctx = Context.auto () in
   let%op x = { x = uniform1 (); o = [ 2; 3 ] } in
@@ -51,7 +51,7 @@ let capture_for_computation () =
   Train.printf ~here:[%here] ~with_code:false ~with_grad:false dim_calc
 
 let test_set_dim_and_set_equal () =
-  let open Operation.DSL_modules in
+  let open Nn_blocks.DSL_modules in
   Tensor.unsafe_reinitialize ();
   Stdio.printf "\n=== Testing set_dim and set_equal functionality ===\n";
 
@@ -139,7 +139,7 @@ let test_set_dim_and_set_equal () =
   Stdio.printf "=== All tests completed ===\n"
 
 let capture_for_shape_validation () =
-  let open Operation.DSL_modules in
+  let open Nn_blocks.DSL_modules in
   Tensor.unsafe_reinitialize ();
   Shape.unsafe_reinitialize ();
   let ctx = Context.auto () in
@@ -253,7 +253,7 @@ let capture_for_shape_validation () =
   Stdio.printf "=== Shape inference integration tests completed ===\n"
 
 let capture_for_shape_inference () =
-  let open Operation.DSL_modules in
+  let open Nn_blocks.DSL_modules in
   Tensor.unsafe_reinitialize ();
   Shape.unsafe_reinitialize ();
   let ctx = Context.auto () in
diff --git a/test/einsum/test_surjectivity.ml b/test/einsum/test_surjectivity.ml
@@ -2,7 +2,7 @@ open Base
 open Ocannl
 module IDX = Train.IDX
 module CDSL = Train.CDSL
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 
 module type Backend = Ir.Backend_intf.Backend
 
diff --git a/test/operations/hello_world_op.ml b/test/operations/hello_world_op.ml
@@ -3,7 +3,7 @@ open Ocannl
 module Tn = Ir.Tnode
 module IDX = Train.IDX
 module CDSL = Train.CDSL
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 
 module type Backend = Ir.Backend_intf.Backend
 
diff --git a/test/operations/micrograd_demo_logging.ml b/test/operations/micrograd_demo_logging.ml
@@ -1,7 +1,7 @@
 open Base
 open Ocannl
 module IDX = Train.IDX
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 module CDSL = Train.CDSL
 module Asgns = Ir.Assignments
 module Tn = Ir.Tnode
diff --git a/test/operations/primitive_ops.ml b/test/operations/primitive_ops.ml
@@ -3,7 +3,7 @@ open Ocannl
 module Tn = Ir.Tnode
 module IDX = Train.IDX
 module CDSL = Train.CDSL
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 
 module type Backend = Ir.Backend_intf.Backend
 
diff --git a/test/operations/test_extend_dsls.ml b/test/operations/test_extend_dsls.ml
@@ -1,5 +1,5 @@
 open! Base
-open Ocannl.Operation.DSL_modules
+open Ocannl.Nn_blocks.DSL_modules
 
 (* Test module-level expansion *)
 [%%extend_dsls
diff --git a/test/operations/test_record_syntax.ml b/test/operations/test_record_syntax.ml
@@ -1,4 +1,4 @@
-open Ocannl.Operation.DSL_modules
+open Ocannl.Nn_blocks.DSL_modules
 
 (* Test %op record syntax with different initialization patterns *)
 
diff --git a/test/operations/test_threefry4x32.ml b/test/operations/test_threefry4x32.ml
@@ -2,7 +2,7 @@ open Base
 module Tn = Ir.Tnode
 module Asgns = Ir.Assignments
 module Ndarray = Ir.Ndarray
-open Ocannl.Operation.DSL_modules
+open Ocannl.Nn_blocks.DSL_modules
 module O = TDSL.O
 
 let%expect_test "threefry4x32 basic test" =
diff --git a/test/operations/test_threefry_precision.ml b/test/operations/test_threefry_precision.ml
@@ -4,7 +4,7 @@ open Ocannl
 let () =
   Utils.settings.output_debug_files_in_build_directory <- true;
   Utils.settings.log_level <- 1;
-  let open Operation.DSL_modules in
+  let open Nn_blocks.DSL_modules in
   (* Create a simple Threefry4x32 operation *)
   let key = TDSL.number ~label:[ "key" ] 42.0 in
   let counter = TDSL.number ~label:[ "counter" ] 1.0 in
diff --git a/test/operations/test_uniform1.ml b/test/operations/test_uniform1.ml
@@ -2,7 +2,7 @@ open Base
 module Tn = Ir.Tnode
 module Asgns = Ir.Assignments
 module Ndarray = Ir.Ndarray
-open Ocannl.Operation.DSL_modules
+open Ocannl.Nn_blocks.DSL_modules
 
 let uniform1_basic_test () =
   Tensor.unsafe_reinitialize ();
@@ -63,7 +63,7 @@ let uniform1_basic_test () =
 let uniform_at1_test () =
   Tensor.unsafe_reinitialize ();
   let ctx = Context.auto () in
-  let open Ocannl.Operation.DSL_modules in
+  let open Ocannl.Nn_blocks.DSL_modules in
   let module O = TDSL.O in
   (* Create a counter tensor *)
   let counter = TDSL.range 5 in
@@ -91,7 +91,7 @@ let uniform_at1_test () =
 let uniform1_shape_preservation_test () =
   Tensor.unsafe_reinitialize ();
   let ctx = Context.auto () in
-  let open Ocannl.Operation.DSL_modules in
+  let open Ocannl.Nn_blocks.DSL_modules in
   let module O = TDSL.O in
   (* Test that uniform1 preserves shape (1:1 mapping) *)
   let input_size = 24 in
diff --git a/test/operations/test_vec_simple.ml b/test/operations/test_vec_simple.ml
@@ -1,5 +1,5 @@
 open Base
-open Ocannl.Operation.DSL_modules
+open Ocannl.Nn_blocks.DSL_modules
 
 let () =
   let ctx = Context.auto () in
diff --git a/test/operations/test_where_precision.ml b/test/operations/test_where_precision.ml
@@ -1,6 +1,6 @@
 open Base
 module Train = Ocannl.Train
-open Ocannl.Operation.DSL_modules
+open Ocannl.Nn_blocks.DSL_modules
 module Tn = Ir.Tnode
 
 let () =
diff --git a/test/operations/test_where_simple.ml b/test/operations/test_where_simple.ml
@@ -1,7 +1,7 @@
 open Base
 open Ocannl
 open Operation.At
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 
 let () =
   Tensor.unsafe_reinitialize ();
diff --git a/test/operations/threefry4x32_demo.ml b/test/operations/threefry4x32_demo.ml
@@ -1,7 +1,7 @@
 (* A simple demo of using the Threefry4x32 PRNG in OCANNL *)
 
 open Base
-open Ocannl.Operation.DSL_modules
+open Ocannl.Nn_blocks.DSL_modules
 module O = TDSL.O
 
 let () =
diff --git a/test/operations/top_down_prec.ml b/test/operations/top_down_prec.ml
@@ -2,7 +2,7 @@
 
 open Base
 module Train = Ocannl.Train
-open Ocannl.Operation.DSL_modules
+open Ocannl.Nn_blocks.DSL_modules
 module Tn = Ir.Tnode
 
 let () =
diff --git a/test/operations/transformer_test.ml b/test/operations/transformer_test.ml
@@ -1,5 +1,5 @@
 open! Base
-open Ocannl.Operation.DSL_modules
+open Ocannl.Nn_blocks.DSL_modules
 
 let () =
   (* Basic transformer test *)
diff --git a/test/operations/zero2hero_1of7.ml b/test/operations/zero2hero_1of7.ml
@@ -4,7 +4,7 @@ module Nd = Ir.Ndarray
 module Asgns = Ir.Assignments
 module IDX = Train.IDX
 module CDSL = Train.CDSL
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 
 module type Backend = Ir.Backend_intf.Backend
 
diff --git a/test/ppx/test_ppx_name_conflict.ml b/test/ppx/test_ppx_name_conflict.ml
@@ -1,5 +1,5 @@
 open Ocannl
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 
 let%op test_inline_defs x =
   let q = { w } * x in
diff --git a/test/ppx/test_ppx_name_conflict_expected.ml b/test/ppx/test_ppx_name_conflict_expected.ml
@@ -1,5 +1,5 @@
 open Ocannl
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 let test_inline_defs =
   let w =
     [%ocaml.error
diff --git a/test/ppx/test_ppx_op.ml b/test/ppx/test_ppx_op.ml
@@ -1,6 +1,6 @@
 open Base
 open Ocannl
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 
 let%op y0 = (2 *. { hey1 }) + 3
 let%op y1 x = ({ hey2 } * 2) + x
diff --git a/test/ppx/test_ppx_op_expected.ml b/test/ppx/test_ppx_op_expected.ml
@@ -1,6 +1,6 @@
 open Base
 open Ocannl
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 let y0 =
   let hey1 =
     (TDSL.param ?more_label:None ?value:None ?values:None ?param_init:None
diff --git a/test/training/bigram.ml b/test/training/bigram.ml
@@ -4,7 +4,7 @@ open Stdio
 open Bigarray
 module Tn = Ir.Tnode
 module IDX = Train.IDX
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 module CDSL = Train.CDSL
 module Asgns = Ir.Assignments
 
diff --git a/test/training/bigram_mlp.ml b/test/training/bigram_mlp.ml
@@ -4,7 +4,7 @@ open Stdio
 open Bigarray
 module Tn = Ir.Tnode
 module IDX = Train.IDX
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 module CDSL = Train.CDSL
 module Asgns = Ir.Assignments
 
diff --git a/test/training/moons_demo.ml b/test/training/moons_demo.ml
@@ -2,7 +2,7 @@ open Base
 open Ocannl
 module Tn = Ir.Tnode
 module IDX = Train.IDX
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 module Asgns = Ir.Assignments
 
 let main () =
diff --git a/test/training/moons_demo_mini.ml b/test/training/moons_demo_mini.ml
@@ -2,7 +2,7 @@ open Base
 open Ocannl
 module Tn = Ir.Tnode
 module IDX = Train.IDX
-open Operation.DSL_modules
+open Nn_blocks.DSL_modules
 module Asgns = Ir.Assignments
 
 let main () =
