gspr:

By my (apparently wrong) understanding, backward should take a Tensor A of shape [m, 3] (for some m depending on the function following f in the composition) and return a Tensor of shape [m, 2] representing the gradient of the composition at some point (obtained from the context object).


No, the backward takes in a tensor of shape [3], representing dL/d output for some scalar L, and returns a tensor of shape [2], representing [dL/ dx, dL/dy].