Add more features to the pre-segmentation pass to move take_along_axis ahead of certain unary ops

The PR add a pre-segmentation pass that moved a take_along_axis operator ahead of certain unary pointwise operators such as castOp and the squeeze op.

The implementation had certain limitations. We should extend the current implementation to remove these.

• We only consider fusions with a single take_along_axis op. We should extend our implementation to support the case where there are multiple take_along_axis ops.
• We move a take_along_axis ahead of a squeeze op. Currently we do so only if a single dimension is being squeezed. We should remove this limitation. This will mean that we have to add multiple unsqueeze ops to the indexTv of the take_along_axis.