You should be able to replace the for loop in j by replacing matmul with broadcasting (or einsum if you can’t help it). Then you can just write b = b + X, eliminating the inplace update.
Best regards
Thomas