Function predict_with_output_summed

intermediate_source/jacobians_hessians.py:290–291

(weight, bias, x)

# and then computing the Jacobian of that function:

def predict_with_output_summed(weight, bias, x):
    return predict(weight, bias, x).sum(0)

batch_jacobian1 = jacrev(predict_with_output_summed, argnums=2)(weight, bias, x).movedim(1, 0)
assert torch.allclose(batch_jacobian0, batch_jacobian1)