Fix tensor shapes for elementwise binary operations with broadcasting

[soumyac1999]

Moved from #738

Description of changes:

Example: [16, 1, 2] + [16, 1552, 1] has output shape of [16, 1552, 2] but the existing code was setting the size to [16, 1, 2].

Related Issues:

Linked Issues:

• Issue #

Issues closed by this PR:

• Closes Investigate possible issues in output shape computation in keras/layers/merge.py #617

Before merging:

• Did you update the flexflow-third-party repo, if modifying any of the Cmake files, the build configs, or the submodules?

---

This change is 

[soumyac1999]

Hi @jiazhihao, could you please re-review this PR? You had already reviewed it in #738 but I wasn't able to get the CI working there so created this new PR.