let where ?(label = []) ~grad_spec t1 t2 t3 =

let module NTDSL = NTDSL_before_div in

let%cd op_asn ~v ~t1 ~t2 ~t3 ~projections = v =: where v1 v2 v3 in

let%cd grad_asn ~t:_ ~g ~t1 ~t2 ~t3 ~projections =

g2 =+ where v1 g 0;

in

Tensor.ternop ~label:("where" :: label) ~ternary_op:Pointwise_tern ~op_asn ~grad_asn ~grad_spec t1

t2 t3

let recip_sqrt = recip_sqrt ~grad_spec:If_needed

let tanh = tanh ~grad_spec:If_needed

let where = where ~grad_spec:If_needed

module NDO = struct

let recip_sqrt = recip_sqrt ~grad_spec:Prohibit_grad

let tanh = tanh ~grad_spec:Prohibit_grad

let where = where ~grad_spec:Prohibit_grad

module TDSL = struct

}

| { pexp_desc = Pexp_ident { txt = Lident op_ident; _ }; _ } when is_primitive_op op_ident ->

default_result

| [%expr [%e? expr1] **. [%e? { pexp_desc = Pexp_constant (Pconst_integer _); _ } as i]] ->

(* FIXME: `**.` should take a tensor and require that it's a literal. *)

(* We need to hardcode these two patterns to prevent the numbers from being converted to tensors. *)

In the first case, we have a binary assignment calculated pointwise. The resulting representation is `Accum_binop` where `accum` is `Add` and `op` is `Mul` (multiplication). In the second case, `*.` is not recognized as one of the built-in operators. This leaves the expression `learning_rate *. p.grad` un-transformed. Since `(*.)` is bound in `NTDSL.O` to pointwise tensor multiplication, this creates an intermediate tensor, that is then added onto p. The resulting representation is `Accum_unop` where `accum` is `Add` and `op` is `Identity`. Both variants end up with the same result, and even with the same computation, because the second variant's computation will get optimized (unless configured not to).