let jacobian f args = let n = Array.length args in let tags = Array.init n (fun _ -> new_tag ()) in let dargs = Array.mapi (fun i arg -> D(tags.(i), arg, one)) args in let result = f dargs in Array.init n (fun i -> let fi = result.(i) in Array.init n (fun j -> let tj = tags.(j) in Array.fold_left (fun res tag -> if tag = tj then extract_deriv tag res else drop_deriv tag res) fi tags))That's efficient in the sense that it doesn't re-compute results within f. Unfortunately, it introduces n tags, where n is the dimensionality of the argument to f. Each tag labels a fundamental ``incremental value'' dx, dy, .... A particular D(...) expression is a tree which represents an expression a0 + a1*dx + a2*dx*dy + ... which potentially contains one term for each of the 2^n different combinations of the dx, dy, .... This means that, though we save on repeated invocations of f, the cost for computing the jacobian becomes exponential in the dimension of the argument because each operation (+, sin, etc) has to manipulate these trees to a depth of 2^n.
The following simple change wastes effort computing f dargs multiple times, but scales only linearly in complexity with the dimension of args. This is a huge win. I've also added a function insert_tag which inserts a new D(tag, x, one) term within the derivative-computing wrappers (but does so in way consistent with the heap property of the diff-tree). The modified functions are below:
let rec insert_tag tag = function | (C(x) as d) -> D(tag, d, one) | (D(tagd, x, dx) as d) when tag < tagd -> D(tag, d, one) | D(tagd, x, dx) when tagd = tag -> raise (Failure "Cannot insert tag which is already present.") | D(tagd, x, dx) -> D(tagd, insert_tag tag x, dx) let d f = fun x -> let tag = new_tag () in let res = f (insert_tag tag x) in extract_deriv tag res let jacobian f args = let n = Array.length args in let jac = Array.make_matrix n n zero in for j = 0 to Pervasives.(-) n 1 do let tag = new_tag () in let dargs = Array.mapi (fun i arg -> if Pervasives.(=) i j then insert_tag tag arg else arg) args in let result = f dargs in Array.iteri (fun i res -> jac.(i).(j) <- extract_deriv tag res) result done; jac