Skip to content

Fix torch.randint generating values equal to the upper bound (#2568)#2714

Merged
TobyRoseman merged 2 commits into
apple:mainfrom
john-rocky:fix-randint-upper-bound
May 27, 2026
Merged

Fix torch.randint generating values equal to the upper bound (#2568)#2714
TobyRoseman merged 2 commits into
apple:mainfrom
john-rocky:fix-randint-upper-bound

Conversation

@john-rocky
Copy link
Copy Markdown
Contributor

Summary

torch.randint(low, high, ...) draws integers in [low, high) — the upper bound is exclusive. The lowering used mb.random_uniform (whose bounds are inclusive) and cast the result to int, so the converted model could emit exactly high (issue #2568).

Clamp the cast result to high - 1 with a minimum op so the converted graph matches torch's exclusive upper bound:

-    rand_int = mb.cast(x=rand_uniform, dtype="int32", name=node.name)
+    rand_int_unclipped = mb.cast(x=rand_uniform, dtype="int32")
+    high_int = mb.cast(x=high, dtype="int32")
+    high_minus_one = mb.sub(x=high_int, y=1)
+    rand_int = mb.minimum(x=rand_int_unclipped, y=high_minus_one, name=node.name)

Fixes #2568.

Test plan

New TestRandint::test_upper_bound_is_exclusive (TorchScript + TorchExport frontends; ExecuTorch is skipped because aten.randint.low is not Aten Canonical) asserts the lowered graph wires a minimum clamp into the randint path, so the converted graph cannot produce values >= high.

Verified locally on macOS that the clamp is present in the lowered graph (convert_to="milinternal"):

[torch.export] random_uniform: True | minimum: True

The full test converts to a serialized mlprogram, so it runs end-to-end in CI.

torch.randint is documented as returning integers in [low, high), but the
lowering used mb.random_uniform whose upper bound is inclusive — and
cast-to-int truncates toward zero, so a sample exactly at high stays at
high after casting. The reporter saw torch.randint(0, 100, ...) produce
values up to 100.

Clamp the cast result to high - 1 with mb.minimum so the converted graph
matches torch's exclusive upper bound. Adds a regression test that asserts
the lowered MIL graph contains the clamp op.
model = TestModel().eval()
x = torch.zeros((1000, 100), dtype=torch.int64)
torch_model = export_torch_model_to_frontend(model, (x,), frontend)
inputs = (
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this statement necessary? If so, why?

)
mlmodel = ct.convert(torch_model, inputs=inputs)

# Structural assertion: the lowering must wire a `minimum` clamp into
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't understand what the rest of this test is doing. Why not just get the predictions from the converted model and make sure the max element is 100?

@john-rocky
Copy link
Copy Markdown
Contributor Author

Done — replaced the MIL op-type walk with a prediction check.

  • On the inputs line: it's needed because the TorchScript frontend requires explicit input types, while exported programs already carry them (same pattern as test_tuple_input just above). Added a short comment to that effect.
  • On checking the output: I assert prediction.max() < 100. Since randint's upper bound is exclusive, a correct model tops out at 99; before the clamp, a 100k-element draw hits exactly 100 — so it's < 100 rather than == 100.

Verified locally: passes on this branch, and fails on main (the unclamped cast emits 100), so it catches the regression.

@TobyRoseman
Copy link
Copy Markdown
Collaborator

CI: https://gitlab.com/coremltools1/coremltools/-/pipelines/2553884155

If this passes, I will merge.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

torch.randint generates out of range values in CoreML

2 participants