Thanks a lot for your reply, @soulitzer. So as stated by you, I should find the jacobian of each of  the model’s parameters and A which will lead to a tuple of Jacobian matrices . Actually further on I am required to update A using this calculated gradient. The gradient matrix should be of size equal to A’s shape. How am I supposed to obtain the gradient matrix of required dimensions using the tuple of Jacobians?
Also, on a side note, could you please explain the difference between ‘entire Jacobian’ and ‘linear combination of its rows’. Thank you.
I am using this 1 code to calculate the jacobian.