numpy conditional operation between 2d and 3d array

Question

I am currently working on a project in Python 3.6 with numpy version 1.19.4. I have a piece of code which worked in Numpy version 1.18.1. The code is supposed to take an gra