CPU and MPS differences for argmax

[ghost]

'$argmax' and 'torch_argmax' seem to fail when run on MPS device

Example:

set.seed(100)

a = torch_randn(4, 4)$to(device = "cpu")
a

Output:
torch_tensor
0.3711 1.9380 -0.7856 -1.0108
-0.7952 -0.5836 -0.1242 0.1746
0.9352 0.4852 -1.8873 -0.8021
-1.3524 0.0546 0.6454 0.1427
[ CPUFloatType{4,4} ]

b = torch_randn(4, 4)$to(device = "mps")
b

Output:
torch_tensor
-1.1729 0.5144 1.0309 0.7047
-0.2580 -1.7188 -2.0322 -1.0363
1.4644 -0.7048 1.1910 -0.1933
0.5761 -0.3971 -0.8470 1.5876
[ MPSFloatType{4,4} ]

as.numeric(a$argmax(dim = 2)$to(device = "cpu"))
Output: [1] 2 4 1 3

as.numeric(a$argmax(dim = 2)$to(device = "mps"))
Output: [1] -29077583 50331649 0 0

as.numeric(b$argmax(dim = 2)$to(device = "cpu"))
Output: [1] 3 1 1 4

as.numeric(b$argmax(dim = 2)$to(device = "mps"))
Output: [1] -29077583 50331649 0 0

(On MacOSX Ventura 13.4.1 with R version 4.3.1 (2023-06-16) -- "Beagle Scouts")

[ghost]

Ok, the issue sits with the conversion function as.numeric(), which can only work on CPU located tensors.

[dfalbel]

Thanks for reporting! We should make sure to return an error instead of wrong results