AutoDiffXd Operations in Python May Not Preserve Derivative Shape

[cohnt]

What happened?

Certain operations in pydrake.math, including max and min, may end up changing the number of derivatives an AutoDiffXd variable has. This can become a major issue down the line -- for example, when jacobian on a function, it expects the number of derivatives to be the same for all output values, in order to construct the Jacobian matrix. Here's a simple example of code that exhibits such a problem:

import numpy as np
from pydrake.all import AutoDiffXd, jacobian
import pydrake.math

def func(q):
	# q should be of length 2
	return np.array([q[0], pydrake.math.min(q[1], 1)])

x = np.array([0, 0])
jacobian(func, x)

y = np.array([0, 2])
jacobian(func, y)

The first jacobian computation works as expected, whereas the second throws an error due to the derivative shapes not lining up.

Here's an example in general of derivative shapes changing when calling min and max, or not being checked:

import numpy as np
from pydrake.all import AutoDiffXd
import pydrake.math

x = AutoDiffXd(0, np.zeros(2))
y = AutoDiffXd(1, np.zeros(3))

print(pydrake.math.min(x, y))
print(pydrake.math.max(x, y))

Version

No response

What operating system are you using?

No response

What installation option are you using?

No response

Relevant log output

No response

[jwnimmer-tri]

My understanding of the architectural premise of AutoDiffXd is that a value with empty derivatives, when combined with a value that has non-empty derivatives, is to be treated as-if it had the matching number of implied zeros in its derivative.

Therefore, binary (two-argument) functions that operate on AutoDiffXd are not required to return a value with a non-empty derivatives filled with zeros; they are allowed to return an empty derivatives vector instead, especially in case one of their arguments already had empty derivatives.

Under that premise, the bug is in pydrake.forwarddiff.jacobian for not handling this case.

Does that sound right?

[cohnt]

I'm far from being an expert in the architecture behind AutoDiffXd, but under that premise, I would agree that the bug is in pydrake.forwarddiff.jacobian.

It does leave open the question of how such operations should handle two inputs with different numbers of derivatives, such as in this case:

import numpy as np
from pydrake.all import AutoDiffXd
import pydrake.math

x = AutoDiffXd(0, np.zeros(2))
y = AutoDiffXd(1, np.zeros(3))

print(pydrake.math.min(x, y))
print(pydrake.math.max(x, y))

Is that the type of thing that should throw an error? I can also make a separate issue for that case if you want to separate it out.

[jwnimmer-tri]

Is that the type of thing that should throw an error?

It would be acceptable to throw an error. The precondition of AutoDiffXd binops is that the derivative sizes must be congruent (i.e., when both sizes are >0, then they must be identical.)

At the moment, checking that precondition for min(a, b) or operator*(a, b) is expensive enough that we only do it in Debug builds: https://github.com/RobotLocomotion/drake/blob/v1.17.0/common/autodiffxd.h#L533-L536 This is similar to other places in C++ (e.g., matrix multiplication of doubles) where Eigen's matching-dimension checks (e.g., on the inner product) only happen in Debug builds. This is somewhat foreign to Python users, where numpy provides guaranteed errors on dimension mismatches. Another part of the problem here is that a Matrix<AutoDiffXd> can embody a different derivative size for each individual element, it's not like a numpy or a tensor where it's dense, 3-dimensional storage.

In cases (like min) where the dimension mistake will propagate to the return value and probably be detected later anyway, I am OK with this status quo. For mismatches where the program would experience undefined behavior or segfaults, we should probably look into improving the code to fail-fast instead. My pending implementation of AutoDiff (#17492) will probably address most of those, so I don't think we need a new issue.

[jwnimmer-tri]

I realized there is a more direct way to explain this situation.

Calculation operations like addition always produce new values. So a + b should always return a specific number of derivatives, even if one of a or b was a constant that had no derivatives. Ditto for sum(values) which calculates the fold of the values over addition.

In contrast, operations like https://docs.python.org/3/library/functions.html#min are typically viewed as selection functions, not calculation functions. This is true for both C++ (both min and min_element operate in terms of references or iterators) as well as Python (see docs). In neither case is there any copying.

For example:

>>> import numpy as np
>>> x = np.array([1.0])
>>> y = np.array([2.0])
>>> id(x)
140330071076144
>>> id(y)
140330071074224
>>> min(x, y)
array([1.])
>>> id(min(x, y))
140330071076144    # N.B. Same as id(x) above.

Since there is only selection, no copying, if the minimal value is a constant (no derivatives) it's impossible to set its size in the return value. Therefore, any code that selects from a list of autodiff values will itself need to implement the derivative-size propagation logic, one way or another.

I would agree that the bug is in pydrake.forwarddiff.jacobian.

We'd welcome a PR with the fix, if you'd like to offer. It seems like it would be fairly a narrowly isolated change within that function.

Note that #19623 is also changing some nearby code. Any merge/rebase conflicts would be minor, but it might be worth waiting for that to land first.

[jwnimmer-tri]

BTW I also just realized that the other functions in the same file (derivative and gradient) might have the same bug, of assuming the derivatives size of what the function returns.

[cohnt]

As best I can tell, only pydrake.forwarddiff.jacobian has this bug, since it's using numpy.vstack on the derivatives (which doesn't know to replace empty arrays with arrays of zeros). Since derivative and gradient don't have to collect the derivatives of multiple objects into a matrix, they avoid the problem.

The change required shouldn't conflict with #19623, since the only change required is in the return statement. I'll go ahead and PR a fix.