From 67e28af08f5df585c59a64a34d47817502bedbbf Mon Sep 17 00:00:00 2001
From: LeSingh1 <sshaurya914@gmail.com>
Date: Sun, 9 Aug 2026 15:54:58 -0700
Subject: [PATCH] Read torch rand's dtype from the argument that holds it

aten::rand is rand(size, *, dtype, layout, device, pin_memory), so the
dtype sits right after the size. The converter unpacked

    shape, _, dtype, _, _ = _get_inputs(context, node)

which skips the dtype and reads the layout instead. torch always leaves
the layout as None here, so the requested dtype was unreachable and every
torch.rand produced fp32, whatever was asked for.

The lookup was wrong for the same reason: dtype.val is already a torch
dtype number, so NUM_TO_DTYPE_STRING takes it directly, without the
TORCH_DTYPE_TO_NUM step that only made sense for a torch.dtype object.
randn, right below, already reads inputs[1].

The existing test_rand is parametrized over float16 / float32 / float64
but only checks that the samples fall in [0, 1), which holds either way.
---
 .../converters/mil/frontend/torch/ops.py      |  7 +++--
 .../mil/frontend/torch/test/test_torch_ops.py | 29 +++++++++++++++++++
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/coremltools/converters/mil/frontend/torch/ops.py b/coremltools/converters/mil/frontend/torch/ops.py
index 7f747baf4..1f16c2bc2 100644
--- a/coremltools/converters/mil/frontend/torch/ops.py
+++ b/coremltools/converters/mil/frontend/torch/ops.py
@@ -5724,8 +5724,11 @@ def _parse_positional_args(context, node) -> Tuple[Var]:
 
 @register_torch_op
 def rand(context, node):
-    shape, _, dtype, _, _ = _get_inputs(context, node)
-    dtype = NUM_TO_DTYPE_STRING[TORCH_DTYPE_TO_NUM[dtype.val]] if dtype else "fp32"
+    # aten::rand(size, *, dtype, layout, device, pin_memory). The dtype is at
+    # position 1; position 2 is the layout, which torch always leaves as None here,
+    # so reading it made the requested dtype unreachable.
+    shape, dtype, _, _, _ = _get_inputs(context, node)
+    dtype = NUM_TO_DTYPE_STRING[dtype.val] if dtype is not None and dtype.val is not None else "fp32"
     low, high = mb.cast(x=0.0, dtype=dtype), mb.cast(x=1.0, dtype=dtype)
     rand_uniform = mb.random_uniform(shape=shape, low=low, high=high)
     context.add(rand_uniform, node.name)
diff --git a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py
index 58860798b..8d59a6106 100644
--- a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py
+++ b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py
@@ -4912,6 +4912,35 @@ def forward(self, x):
 
         self.run_compare_torch(shape, TestModel(), backend=backend, compute_unit=compute_unit)
 
+    @pytest.mark.parametrize(
+        "dtype, expected_dtype",
+        [
+            (None, "fp32"),
+            (torch.float16, "fp16"),
+            (torch.float32, "fp32"),
+            (torch.float64, "fp32"),
+        ],
+    )
+    def test_rand_dtype(self, dtype, expected_dtype):
+        """The requested dtype is the argument right after the size in aten::rand."""
+
+        class TestModel(nn.Module):
+            def forward(self, x):
+                y = torch.rand((2, 3)) if dtype is None else torch.rand((2, 3), dtype=dtype)
+                return x + y.to(x.dtype)
+
+        traced = torch.jit.trace(TestModel().eval(), torch.rand(2, 3))
+        prog = ct.convert(
+            traced,
+            inputs=[ct.TensorType(name="x", shape=(2, 3))],
+            convert_to="milinternal",
+            minimum_deployment_target=ct.target.iOS17,
+            compute_precision=ct.precision.FLOAT32,
+        )
+        random_uniform_ops = prog.functions["main"].find_ops(op_type="random_uniform")
+        assert len(random_uniform_ops) == 1
+        assert types.builtin_to_string(random_uniform_ops[0].outputs[0].dtype) == expected_dtype
+
 
 class TestRandLike(TorchBaseTest):
     @pytest.mark.parametrize(