Open martinResearch opened 5 years ago
The assumption of sparsegrad is that each derivative, at every step, is a matrix. sparsegrad is most useful if the final result is a sparse matrix which cannot be represented as a dense matrix (for example, because it has dimension of 1M times 1M elements). Because matrix is indexed by two indices, it fits well to express derivatives of 1-D arrays w.r.t. 1-D arrays. Then one index refers to one array, and the second index refers to the other array: it's Jacobian of a vector valued function w.r.t. a vector valued function. It is relatively easy for sparsegrad to support Nd arrays by using 1-D indexing for Nd arrays. This could be done internally. Let me illustrate the idea by rewriting your example in the following (running) way:
import numpy as np
import sparsegrad.forward as ad
m,n=3,3,
def f(x_flat):
return x_flat-x_flat[indexing[::-1,::-1].flatten()]
indexing = np.arange(m*n).reshape(m,n)
assert (indexing.flatten()==np.arange(m*n)).all()
x=np.zeros((3,3))
x_flat=x.flatten()
# Now, x[i,j] = x_flat[indexing[i,j]]
print(f(ad.seed(x_flat)).dvalue)
# Recovery of derivative: dx[i,j]/dx[m,n] = dvalue[indexing[i,j],indexing[m,n]]
It is not difficult to add such functionality which would allow to run your example unmodified. I think, the most of what is necessary is to:
mshape
broadcast
diag
The bigger question is maybe about the performance and if sprasegrad
is a good library for such a problem.
The design constraint was: derivatives are assumed to be sparse, I wanted pure Python and (as far as I know) scipy only has 2-D sparse matrices. For problems for which dense tensors would be acceptable for derivatives, implementing forward mode AD would be more efficient and simpler with dense arrays. Then you could use ideas from sparsegrad regarding integration with numpy etc.
I've just created a branch ndim
which passes the tests (under Python 3.6) and runs your example. However, I could not check the correctness for n-D>1. Also, this it is far from providing full support of n-D arrays, which would require supporting axis=keyword argument, and wrapping functions which were not useful for 1-D arrays (reshape, transpose, etc.). Also, there could be numpy indexing magic which could improve performance.
However, at this stage, n-D indexing, broadcasting and arithmetics should be supported.
Please let me know if you would like to contribute to merging this functionality into master, for example by testing the library.
Hi, indeed jacobian of functions that take ND arrays as input and ouput ND Arrays are not sparse 2D matrices but (sparse) tensors. We could have a sparse tensor class that internally use a sparse matrix to store the data. Instead, The convention i used in here was to represent the jacobian of the function f by the 2D jacobian of the modified function f_flat(x_flat)= f( np.reshape(xflat,x.shape) ).flatten(). I'll try to find some time soon to help in testing what you have done
It would be good if the API of sparsegrad what following closely the one from autograd, so that it becomes very easy to switch between the two libraries. A large part of the test in the test subfolder could then be directly reused.
This library does not support matrix operations For example this does not work
Would it be possible to add support for operations on matrices and Nd arrays ? For the reference here is a matlab implementation for forward differentiation that support nd arrays.