> Or do you do a for loop to avoid the allocation of an intermediate array with shape of g_x?

exactly, this is for saving memory. Thanks for the proposed improvement, sounds reasonable to me!