an alternative but inefficient way of calculating grad_h would be:

Yes exactly.

Also, in this example, we are processing a batch of 64 examples together. h is of size [64, 100] and so is h_relu . So how do we calculate the Jacobian of matrices?

From the autograd point of view, “batch” does not exist, so I usually find it easier to just consider it a function of 2D -> 2D.
If you want to reuse the logic above, you need to flatten input/output to make then 1D/1D.
And compute the corresponding Jacobian. Since elements in the batch don’t interact accross batches, then the full Jacobian will be a block diagonal matrix. Where each block is the Jacobian you defined for a single sample.