Replace this section with `np.broadcast_shapes()` ? <img width="1780" height="1640" alt="Image" src="https://github.com/user-attachments/assets/6f9d8cca-c4f3-4b2c-affb-68fcf733a18b" />
