如何在自动grad中逐个元素分配方程式

Question

我正在尝试为非线性偏微分方程实现一个基于自动grad的求解器。与大多数偏微分方程一样，我需要能够在输入向量的单个条目中进行操作，但显然这会破坏自动grad。我创建了这个简单的例子来展示我面临的问题：

以下代码有效：

def my_equation(x):
eq = x
return eq
x = np.random.randn(2,)
jac = autograd.jacobian(my_equation)
jacval = jac(x)
print(jacval)

以下代码不起作用：

def my_equation(x):
eq = x
# This breaks the code, although is a 
# trivial line
eq[1] = x[1]
return eq
x = np.random.randn(2,)
jac = autograd.jacobian(my_equation)
jacval = jac(x)
print(jacval)

我在几个地方读到你不能在自动grad中分配元素。这真的是真的吗？有什么解决方法吗？或者也许是另一个图书馆可以推荐？

谢谢！

Answer 1

事实上，数组索引分配在自动grad中是不可能的。人们已经在自动grad中编写了偏微分方程求解器(见 https://github.com/HIPS/autograd/tree/master/examples/fluidsim(，所以也许有一种方法可以在自动grad中解决问题。

JAX 似乎为 jax.ops 包提供了一种解决方法(参见 https://jax.readthedocs.io/en/latest/jax.ops.html 和 https://github.com/google/jax#current-gotchas(。

看来数组索引在 PyTorch 中是可能的。这表明PyTorch将是你的出路。以下代码有效。

import torch
def f(x):
eq = 2*x 
eq[0] = x[0] 
return eq
x = torch.rand(4, requires_grad=True) 
y = f(x)
z = torch.sum(y)
z.backward()
print(x.grad) # prints [1., 2., 2., 2.]

"或者也许是另一个图书馆的建议？">

看看海豚伴随。如果可以用 FEniCS 编写偏微分方程求解器，那么可能会有所帮助。

http://www.dolfin-adjoint.org/en/latest/

请注意，通过非线性方程求解器的朴素反向传播可能是计算标量损失函数导数的低效方法，请参阅 https://cs.stanford.edu/~ambrad/adjoint_tutorial.pdf 以获取有关伴随方法的非常可靠的教程，包括非线性问题。