So, to clarify: suppose I have some function f:R^2 -> R^3. Suppose its formula is f(x,y) = (x^2,xy, y^2). It’s implemented with a forward that takes a single Tensor of shape [2] as argument and returns a single Tensor of shape [3] as per the formula. 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). The gradient of f can be represented by a 3x2 matrix, the Jacobian at the given point, and the value returned from backward is simply the matrix product of A and this matrix.
What’s the point of taking the A argument at all? Or, in your formulation: Why does backward need to know dL/dy, when as you say the result we want is just the product of dL/dy and dy/dx? Why isn’t it just the responsibility of backward to compute dy/dx (at the given point)?