def test_gather_dynamic_overflow_int16(self, overflow):
        """Dynamic input indices should also be cast if x dim size doesn't overflow int16 range."""

        @mb.program(
            input_specs=[
                mb.TensorSpec(shape=(32769 if overflow else 2, 3)),
                mb.TensorSpec(shape=(2,), dtype=types.int32),
            ],
            opset_version=ct.target.iOS17,
        )
        def prog(x, indices):
            return mb.gather(x=x, indices=indices, axis=0)

        prev_prog, _, block = apply_pass_and_basic_check(prog, "common::add_int16_cast")
        if overflow:
            assert get_op_types_in_program(prog) == get_op_types_in_program(prev_prog)
        else:
            assert get_op_types_in_program(prog) == ["cast", "gather"]
            cast_op = block.find_ops(op_type="cast")[0]
            assert cast_op.dtype.val == "int16"
            assert cast_op.outputs[0] == block.find_ops(op_type="gather")[0].indices
Dynamic input indices should also be cast if x dim size doesn't overflow int16 range.
coremltools/converters/mil/mil/passes/tests/test_quantization_passes.py