I'm interested in constructing a NN which has the following structure: NN
I want a sequential NN in which the last layer takes the partials of the previous layer (which contains only 1 activation function) w.r.t. the inputs (so basically, if the input is x, I want the output layer to compute dh(x)/dx, where h is the output of the previous layer). While it is straightfoward to just compute the partials themselves using GradientTape, I'm interested in constructing a separate layer because I want to train the model using backpropagation.