    def test_remove_no_op_concats_exceptions(self):
        """We expect this to raise an exception in these test cases."""

        # AIT expects all concat inputs to have the same rank.
        with self.assertRaises(RuntimeError):
            self._test_remove_no_op_concats_impl(
                input_shapes=[[2, 4], [1]],
                should_keep_concat=False,
                test_name="test_remove_no_op_concats_same_rank",
            )

        # AIT expects all concat inputs to have the same dimension sizes except for the concat_dim.
        with self.assertRaises(RuntimeError):
            self._test_remove_no_op_concats_impl(
                input_shapes=[[2, 4], [0, 0]],
                should_keep_concat=False,
                test_name="test_remove_no_ops_concat_same_dim_sizes",
            )
We expect this to raise an exception in these test cases.
tests/unittest/compiler/test_remove_no_op_concats.py