Bug in find_concat_op

[shivadbhavsar]

Repro:

# concat_err.py
import numpy as np
import migraphx

p = migraphx.program()
mm = p.get_main_module()

lit1 = mm.add_literal(np.random.randn(128, 1, 1, 1, 1).astype(np.float32))
lit2 = mm.add_literal(np.random.randn(1, 3, 1, 1, 1).astype(np.float32))

lit1_bc = mm.add_instruction(migraphx.op('multibroadcast', out_lens=[128, 3, 224, 224, 1]), [lit1])
lit2_bc = mm.add_instruction(migraphx.op('multibroadcast', out_lens=[128, 3, 224, 224, 1]), [lit2])

cat = mm.add_instruction(migraphx.op('concat', axis=4), [lit1_bc, lit2_bc])

mm.add_return([cat])

Compile:
migraphx-driver compile concat_err.py