- [The syntax extension %%extend_dsls](#the-syntax-extension-extend_dsls)

This syntax extension creates a module `DSL_modules` with the same submodules as `Operation.DSL_modules`. It removes the boilerplate associated with introducing new operators into the modules `TDSL`, `NTDSL`, `PDSL` and their `O` submodules. The payload (i.e. content) of `%%extend_dsls` must be non-recursive let-bindings. They are parsed using a slight variant of the `%op` syntax, and are inserted into the DSL modules. One unique feature of `%%extend_dsls` parsing is that inline tensor definitions, like in `%cd`, do not introduce gradients for the tensors, but, like `%op`, they do introduce initialization for the inline-defined tensors.

(modules ppx_shared ppx_cd ppx_op ppx_extend_dsls ppx_ocannl)

open Base

open Ppxlib

open Ppx_shared

let transform_dsl_binding ~loc ~dsl_name binding =

let transform_expr expr =

let vbs, result =

Ppx_op.translate ~no_grads_for_inline_defs:true

@@ add_module_qualifier_to_applied_function ~module_name:dsl_name expr

in

if Map.is_empty vbs then result

else

Ast_builder.Default.pexp_extension ~loc

@@ Location.error_extensionf ~loc

"%%extend_dsls functions with inline definitions must take a unit parameter"

in

let params, pvb_expr =

match binding.pvb_expr with

| { pexp_desc = Pexp_function (params, _, _); _ } as expr -> (params, transform_expr expr)

| _ -> ([], transform_expr binding.pvb_expr)

in

(params, { binding with pvb_expr })

let str_expander ~loc:pstr_loc ~path:_ str_items =

let transform_op_binding params binding =

let loc = binding.pvb_loc in

let label_p =

{ pparam_loc = loc; pparam_desc = Pparam_val (Optional "label", None, [%pat? label]) }

in

let f = function

| { pparam_desc = Pparam_val (label, _, pat); _ } -> (label, pat2expr pat)

| _ -> assert false

in

let args = List.map params ~f in

let body =

Ast_helper.Exp.apply ~loc (pat2expr binding.pvb_pat)

(args @ [ (Optional "label", [%expr label]); (Nolabel, [%expr ()]) ])

in

let pvb_expr =

{

binding.pvb_expr with

pexp_desc = Pexp_function (label_p :: params, None, Pfunction_body body);

}

in

{ binding with pvb_expr }

in

let items_for_dsl dsl_name =

let item_bindings, op_item_bindings =

List.unzip

@@ List.concat_map str_items ~f:(function

| { pstr_desc = Pstr_value (Nonrecursive, bindings); pstr_loc = loc; _ } ->

List.map bindings ~f:(fun binding ->

let params, binding = transform_dsl_binding ~loc ~dsl_name binding in

let op_binding = transform_op_binding params binding in

(binding, op_binding))

| { pstr_loc = loc; _ } ->

let pat = Ast_helper.Pat.var ~loc { txt = "syntax_error"; loc } in

let v =

Ast_builder.Default.pexp_extension ~loc

@@ Location.error_extensionf ~loc

"ppx_extend_dsls: currently only non-recursive value bindings are supported"

in

[ (Ast_helper.Vb.mk ~loc pat v, Ast_helper.Vb.mk ~loc pat v) ])

in

let item = { pstr_desc = Pstr_value (Nonrecursive, item_bindings); pstr_loc } in

let op_item = { pstr_desc = Pstr_value (Nonrecursive, op_item_bindings); pstr_loc } in

(item, op_item)

in

let loc = pstr_loc in

let item_TDSL, op_item_TDSL = items_for_dsl "TDSL" in

let item_NTDSL, op_item_NTDSL = items_for_dsl "NTDSL" in

let item_PDSL, op_item_PDSL = items_for_dsl "PDSL" in

[%stri

module DSL_modules = struct

module Ir = Ir

module Shape = Shape

module Tensor = Tensor

module TDSL = struct

include TDSL

[%%i item_TDSL]

module O = struct

include TDSL.O

[%%i op_item_TDSL]

end

end

module NTDSL = struct

include NTDSL

[%%i item_NTDSL]

module O = struct

include NTDSL.O

[%%i op_item_NTDSL]

end

end

module PDSL = struct

include PDSL

[%%i item_PDSL]

module O = struct

include PDSL.O

[%%i op_item_PDSL]

end

end

end]

Ppxlib.Context_free.Rule.extension

@@ Extension.declare "extend_dsls" Extension.Context.structure_item

Ast_pattern.(pstr __)

Ppx_extend_dsls.str_expander;

(test

(name test_extend_dsls)

(package neural_nets_lib)

(deps

ocannl_config

(env_var OCANNL_BACKEND))

(modules test_extend_dsls)

(libraries base ocannl stdio)

(preprocess

(pps ppx_here ppx_ocannl ppx_expect)))

Retrieving commandline, environment, or config file variable ocannl_log_level

Found 0, in the config file

HERE: test/operations/test_extend_dsls.ml:23:21

[23]: relu shape 0:$80 <not-in-yet>

HERE: test/operations/test_extend_dsls.ml:29:21

[46]: relu shape 0:$161 <not-in-yet>

HERE: test/operations/test_extend_dsls.ml:35:21

[75]: relu shape 0:$242 <not-in-yet>

grad_relu <not-in-yet>

open! Base

open Ocannl.Operation.DSL_modules

[%%extend_dsls

let my_tensor x = TDSL.O.( + ) x !.2.0 + !.1.0

let my_complex_op a b c =

let d = a * b in

let e = d + c in

relu e

let my_random_op () = { u1 = uniform () } - { u2 = uniform () }]

let () =

(* Test that the functions are available in each DSL module *)

let open! DSL_modules.TDSL.O in

let x = uniform () in

let result = my_complex_op (my_tensor x) (my_random_op ()) x in

Tensor.print ~here:[%here] ~force:false ~with_code:false ~with_grad:true `Inline result

let () =

let open! DSL_modules.NTDSL.O in

let x = uniform () in

let result = my_complex_op (my_tensor x) (my_random_op ()) x in

Tensor.print ~here:[%here] ~force:false ~with_code:false ~with_grad:true `Inline result

let () =

let open! DSL_modules.PDSL.O in

let x = uniform () in

let result = my_complex_op (my_tensor x) (my_random_op ()) x in

Tensor.print ~here:[%here] ~force:false ~with_code:false ~with_grad:true `Inline result