Array Differentiabilty Support for Clad

[grimmmyshini]

So, following up on the meeting the other day, just wanted to share some ideas over here (and some findings while I briefly experimented with this) and maybe this could garner other ideas/responses. @ioanaif, this might be helpful to you!

So, Clad already supports differentiating arrays but only as dependant variables. We also have support for independent arrays named "p". Naively removing the identifier checking (hence allowing all <type>*/<type>[] variants to be treated like "p") results in very interesting results; of which a subset I have compiled here with some useful (albeit sparse) comments.

It seems as though we need to create unique identifiers for each array element because currently, you will see that Clad differentiates all array references a[x] given that the subscript is the same as the variable we are differentiating with respect to. Although this is just speculation.

Another issue is thinking of how to return array derivatives. We cannot expect _result to contain all the array diffs as that introduces a lot of boundary ambiguity and scaling issues. A simple solution is to return a mirror of the original pointer/array that points to the differentiated values, eg

double f(double* a1, double* a2, ..., double* an){ ... }

should be differentiated to

void f_grad(double* a1, double* a2, ..., double* _d_an, double* _result){ ... }

Although we can see how easily that gets out of hand. Maybe there is a better way to return array gradients. perhaps a double** _result?

Anyhow, any discussion/improvements/ideas are all welcome here!

[grimmmyshini]

@efremale Suggests:

it makes sense to use the idea with "mirror" inputs for the reverse mode, it is indeed a better way of doing it compared to the current implementation with a flat double* gradient.

Especially for functions like

double f(double x, double* y, double** z)

it is going to be pretty hard to map all the gradients into a single flat vector and

so double f_grad(double x, double& _d_x, double* y, double* _d_y, double**z, double** _d_z) seems like a good idea.

I think there still are some open questions regarding the design and implementation of this approach, for example, with the old idea with flat double* result, the type of the gradient function generated by some calls are as follows:

clad::gradient(f, "x") --> double (double, double*, double**, double* /*flat gradient result*/)
clad::gradient(f, "y") --> double (double, double*, double**, double* /*flat gradient result*/)
clad::gradient(f, "x, y") --> double (double, double*, double**, double* /*flat gradient result*/)
i.e all calls to the same function return the same signature of the gradient function.

With the new idea with mirrors, the type of the gradient function generated by the same calls are as follows:

clad::gradient(f, "x") --> double (double, double& /*_d_x*/, double*, double**)
clad::gradient(f, "y") --> double (double, double*, double* /*_d_y*/, double**) ... etc.
i.e all calls to same functions return different derivative signatures.

That is, given a function f we are now unable to determine the return type of the call to clad::gradient(f, args) before parsing args with Clad (with the old idea we could determine the return type based on the type of f alone and that did not require invoking Clad plugin). This may be a problem since it seems Clang wants to determine the type of everything in the compilation units before all the plugins (e.g. Clad) are invoked. Although there might be workarounds for the same.

[bradbell]

Suppose that one has a function

double f(double* x)
{ ... }

Is there an example that shows how to differentiate such a function ?

[parth-07]

Hi @bradbell ,

Here's an example that describes differentiation of such a function:

#include <iostream>
#include "clad/Differentiator/Differentiator.h"

double fn(double *x) {
    return 5 * x[0] + 7 * x[1];
}

int main() {
    auto f_grad = clad::gradient(fn);
    double x[3] = {1, 3, 5};
    double dx[3] = {0, 0, 0};
    clad::array_ref<double> dx_ref(dx, 3);
    f_grad.execute(x, dx);
    std::cout << "dfn/dx[0]: " << dx[0] << ", dfn/dx[1]: " << dx[1] << ", dfn/dx[2]: " << dx[2] << "\n";
}

This outputs:

dfn/dx[0]: 5, dfn/dx[1]: 7, dfn/dx[2]: 0

This works, but Clad currently has limited support of pointers in the reverse-mode. So, more advanced example that use pointer features such as pointer arithmetic may fail.