Skip to content

Fix and unblock TestConsistency for median #241

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Ronian526 merged 6 commits into from
Jan 27, 2023
Merged

Fix and unblock TestConsistency for median #241

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e0f28d6
Fix and unblock TestConsistency for median
Ronian526 Jan 23, 2023
cbc8d85
- cast output indices to int64
Ronian526 Jan 20, 2023
cb4c34e
- adding test_median_int16 test
Ronian526 Jan 27, 2023
698a710
- rebase
Ronian526 Jan 23, 2023
d5372ad
- remove casting back
Ronian526 Jan 25, 2023
93f7710
- remove casting back
Ronian526 Jan 25, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 27 additions & 4 deletions aten/src/ATen/native/mps/operations/ReduceOps.mm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1847,6 +1847,17 @@ Tensor median_mps(const Tensor& input_t) {
MPSGraphTensor * reshapedTensor = [mpsGraph reshapeTensor:inputTensor
withShape:@[@-1]
name:nil];
MPSDataType dataType = [inputTensor dataType];
// #issue 104398441 sortWithTensor only supports following types, cast if necessary
if (dataType != MPSDataTypeInt32 &&
Copy link

@ssaladis ssaladis Jan 19, 2023

Choose a reason for hiding this comment

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

Maybe you can use a tuple and then compare with it instead of repeating the same conditional here and below ?

dataType != MPSDataTypeFloat32 &&
dataType != MPSDataTypeFloat16) {
dataType = (dataType & MPSDataTypeFloatBit) ? MPSDataTypeFloat32 : MPSDataTypeInt32;
reshapedTensor = [mpsGraph castTensor:reshapedTensor
toType:dataType
name:@"castReshapedTensor"];
}

MPSGraphTensor * sortedTensor = [mpsGraph
sortWithTensor:reshapedTensor
axis:((NSUInteger) (int)0)
Expand Down Expand Up @@ -1934,7 +1945,7 @@ Tensor median_mps(const Tensor& input_t) {
auto stream = at::mps::getCurrentMPSStream();

@autoreleasepool {
string key = func_name + ":" + to_string(dim_) + ":" + native_mps::getTensorsStringKey(input_t);
string key = func_name + ":" + to_string(dim_) + ":" + native_mps::getTensorsStringKey(input_t) + ":" + native_mps::getTensorsStringKey(indices_t);
CachedGraph* cachedGraph = cache_->LookUpAs<CachedGraph>(key);

if(!cachedGraph) {
Expand All @@ -1948,8 +1959,20 @@ Tensor median_mps(const Tensor& input_t) {

MPSGraphTensor* inputTensor = native_mps::mpsGraphUnrankedPlaceHolder(mpsGraph, native_mps::getMPSDataType(input_t.scalar_type()));
MPSGraphTensor* outputTensor = nil;
MPSGraphTensor* castInputTensor = inputTensor;
MPSDataType dataType = native_mps::getMPSDataType(input_t.scalar_type());
// #issue 104398441 sortWithTensor only supports following types, cast if necessary
if (dataType != MPSDataTypeInt32 &&
dataType != MPSDataTypeFloat32 &&
dataType != MPSDataTypeFloat16) {
dataType = (dataType & MPSDataTypeFloatBit) ? MPSDataTypeFloat32 : MPSDataTypeInt32;
castInputTensor = [mpsGraph castTensor:inputTensor
toType:dataType
name:@"castInputTensor"];
}

MPSGraphTensor * sortedTensor = [mpsGraph
sortWithTensor:inputTensor
sortWithTensor:castInputTensor
axis:((NSUInteger) (int)dim_)
name:nil];

Expand All @@ -1959,7 +1982,7 @@ Tensor median_mps(const Tensor& input_t) {
length:1
name:nil];
MPSGraphTensor* argreduceOutTensor = nil;
argreduceOutTensor = [mpsGraph argSortWithTensor:inputTensor
argreduceOutTensor = [mpsGraph argSortWithTensor:castInputTensor
axis:(NSInteger)dim_
name:@"argmax_out"];
MPSGraphTensor* argOutputTensor = [mpsGraph sliceTensor:argreduceOutTensor
Expand Down Expand Up @@ -2038,7 +2061,7 @@ Tensor median_mps(const Tensor& input_t) {
int64_t num_input_dims = input_shape.size();
NSMutableArray<NSNumber*> *apparent_out_shape = nil;
// Use this if keepdim is false
int64_t num_output_dims = num_input_dims - 1;
int64_t num_output_dims = num_input_dims - 1 < 0 ? 0 : num_input_dims - 1;

std::vector<int64_t> vec_apparent_out_shape(num_input_dims);
std::vector<int64_t> vec_out_shape(num_output_dims);
Expand Down
2 changes: 1 addition & 1 deletion aten/src/ATen/native/mps/operations/Unique.mm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@
return @[resultTensor, inverseIndicesTensor, countTensor, lengthTensor];
}

// Sort only supports following types, cast if necessary
// #issue 104398441 sortWithTensor only supports following types, cast if necessary
if (dataType != MPSDataTypeInt32 &&
dataType != MPSDataTypeFloat32 &&
dataType != MPSDataTypeFloat16) {
Expand Down
12 changes: 11 additions & 1 deletion test/test_mps.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2237,6 +2237,17 @@ def helper(shape):
# helper((2, 8, 4, 5), torch.int16)
# helper((2, 8, 4, 5), torch.int32)

def test_median_int16(self):
def helper(shape, dtype):
cpu_x = torch.randint(-9999, 9999, shape, device='cpu', dtype=dtype)
x = cpu_x.detach().clone().to('mps')

median_result = torch.median(x)
median_result_cpu = torch.median(cpu_x)
self.assertEqual(median_result, median_result_cpu)

helper((2, 8, 4, 5), torch.int16)

class TestLogical(TestCase):
def _wrap_tensor(self, x, device="cpu", dtype=None, requires_grad=False):
return torch.tensor(x, device=device, dtype=dtype, requires_grad=requires_grad)
Expand Down Expand Up @@ -9369,7 +9380,6 @@ class TestConsistency(TestCase):
BLOCKLIST = {
# Functions that hard crash
'nn.functional.softplus': [torch.float32],
'median': [torch.float32, torch.int16, torch.int32, torch.uint8, torch.int16],
'sgn': [torch.bool],
'linalg.inv': [torch.float32],
'linalg.inv_ex': [torch.float32],
Expand Down