Add regression test for decoder_only_block/decoder_only API

Exercises the new Nn_blocks.decoder_only helper with a 2-layer stack,
causal mask, and forward pass, validating output shape. This ensures
the new public API added in the previous commit has CI coverage.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: lukstafi

test/operations/decoder_only_test.expected

@@ -0,0 +1,16 @@
+Retrieving commandline, environment, or config file variable ocannl_log_level
+Found 0, in the config file
+Testing decoder_only (2-layer stack)
+Output shape:
+((batch
+  ((dims
+    ((Dim ((d 2) (label ()) (proj_id ((Proj_id 1)))))
+     (Dim ((d 4) (label ()) (proj_id ((Proj_id 2)))))))
+   (bcast Broadcastable) (prov (((sh_id 307) (kind Batch))))))
+ (input
+  ((dims ()) (bcast Broadcastable) (prov (((sh_id 307) (kind Input))))))
+ (output
+  ((dims ((Dim ((d 16) (label ()) (proj_id ((Proj_id 3)))))))
+   (bcast Broadcastable) (prov (((sh_id 307) (kind Output))))))
+ (batch_padding ()) (input_padding ()) (output_padding ()) (padding_elem ())
+ (id 307) (debug_name layer_norm))

test/operations/decoder_only_test.ml

@@ -0,0 +1,38 @@
+open! Base
+open Ocannl.Nn_blocks.DSL_modules
+
+let () =
+  let ctx = Context.auto () in
+  let batch_size = 2 in
+  let seq_len = 4 in
+  let d_model = 16 in
+  let num_heads = 2 in
+  let d_ff = 32 in
+
+  Stdio.printf "Testing decoder_only (2-layer stack)\n";
+
+  (* decoder_only internally creates decoder_only_block instances,
+     so this exercises both functions. *)
+  let stack =
+    Ocannl.Nn_blocks.decoder_only ~label:[ "test_stack" ] ~num_layers:2 ~num_heads ~d_k:d_model
+      ~d_v:d_model ~d_ff ()
+  in
+
+  let input =
+    TDSL.range_of_shape ~label:[ "input" ] ~batch_dims:[ batch_size; seq_len ] ~input_dims:[]
+      ~output_dims:[ d_model ] ()
+  in
+
+  let mask =
+    NTDSL.init ~l:"mask" ~prec:Ir.Ops.single ~b:[ seq_len ] ~i:[ seq_len ] ~o:[]
+      ~f:(function
+        | [| s; t |] -> if s >= t then 1. else 0.
+        | _ -> failwith "unexpected mask indices")
+      ()
+  in
+
+  let output = stack ~train_step:None input ~mask in
+  let _ctx = Ocannl.Train.forward_once ctx output in
+
+  Stdio.printf "Output shape:\n%s\n%!"
+    (Sexp.to_string_hum ([%sexp_of: Shape.t] output.Tensor.shape))

test/operations/dune

@@ -377,6 +377,17 @@
  (preprocess
   (pps ppx_here ppx_ocannl ppx_sexp_conv)))
 
+(test
+ (name decoder_only_test)
+ (package neural_nets_lib)
+ (deps
+  ocannl_config
+  (env_var OCANNL_BACKEND))
+ (modules decoder_only_test)
+ (libraries base ocannl)
+ (preprocess
+  (pps ppx_here ppx_ocannl ppx_sexp_conv)))
+
 (test
  (name test_extend_dsls)
  (package neural_nets_lib)