[MPS] Register norm_dtype_out_mps and cdist (#91643)
Add support for `norm_dtype_out` and `cdist` ops
Pull Request resolved: https://github.com/pytorch/pytorch/pull/91643
Approved by: https://github.com/razarmehr
diff --git a/aten/src/ATen/native/mps/OperationUtils.h b/aten/src/ATen/native/mps/OperationUtils.h
index b0a3389..3066671 100644
--- a/aten/src/ATen/native/mps/OperationUtils.h
+++ b/aten/src/ATen/native/mps/OperationUtils.h
@@ -43,6 +43,8 @@
MPSDataType getMPSScalarType(ScalarType scalar_type);
MPSScalar getMPSScalar(const Scalar& scalar, ScalarType type);
std::string getMPSTypeString(ScalarType scalar_type);
+NSArray<NSNumber*>* getTensorAxes(const Tensor& t);
+NSArray<NSNumber*>* getTensorAxes(const Tensor& t, at::OptionalIntArrayRef dim);
std::string getMPSShapeString(MPSShape* shape);
std::string getTensorsStringKey(const TensorList& tensors, bool use_scalar_value = false);
std::string getArrayRefString(const IntArrayRef s);
@@ -127,6 +129,13 @@
MPSGraphTensor *outputTensor_ = nil;
};
+struct MPSBinaryCachedGraph : public MPSCachedGraph
+{
+ MPSBinaryCachedGraph(MPSGraph *graph) : MPSCachedGraph(graph) {}
+ MPSGraphTensor *inputTensor_ = nil;
+ MPSGraphTensor *otherTensor_ = nil;
+ MPSGraphTensor *outputTensor_ = nil;
+};
// TODO: Improve the overall design of MPSGraphCache.
// https://github.com/pytorch/pytorch/issues/77176
diff --git a/aten/src/ATen/native/mps/OperationUtils.mm b/aten/src/ATen/native/mps/OperationUtils.mm
index 9da3b2e..8de0016 100644
--- a/aten/src/ATen/native/mps/OperationUtils.mm
+++ b/aten/src/ATen/native/mps/OperationUtils.mm
@@ -87,6 +87,30 @@
}
}
+NSArray<NSNumber*>* getTensorAxes(const Tensor& t) {
+ int64_t ndim = t.dim();
+ auto axes = [NSMutableArray<NSNumber*> arrayWithCapacity:ndim];
+ for (const auto i: c10::irange(ndim)) {
+ axes[i] = [NSNumber numberWithInteger:i];
+ }
+ return axes;
+}
+
+NSArray<NSNumber*>* getTensorAxes(const Tensor& t, at::OptionalIntArrayRef dim) {
+ if (dim.has_value() && dim.value().size() != 0) {
+ IntArrayRef dimValues = dim.value();
+ int ndim = dimValues.size();
+ auto axes = [NSMutableArray<NSNumber*> arrayWithCapacity:ndim];
+ for (const auto i: c10::irange(ndim)) {
+ axes[i] = [NSNumber numberWithInteger:dimValues[i]];
+ }
+
+ return axes;
+ }
+
+ return getTensorAxes(t);
+}
+
std::string getMPSShapeString(MPSShape* shape) {
std::string str;
for(NSNumber *elem in shape) {
diff --git a/aten/src/ATen/native/mps/operations/ReduceOps.mm b/aten/src/ATen/native/mps/operations/ReduceOps.mm
index 43b3336..9ef87b4 100644
--- a/aten/src/ATen/native/mps/operations/ReduceOps.mm
+++ b/aten/src/ATen/native/mps/operations/ReduceOps.mm
@@ -15,6 +15,9 @@
namespace at {
namespace native {
+typedef MPSGraphTensor* (^NormOpBlock)(mps::MPSBinaryCachedGraph*, MPSGraphTensor*, MPSGraphTensor*);
+#define NormOpFn(graph, primary, secondary) MPSGraphTensor* (mps::MPSBinaryCachedGraph* graph, MPSGraphTensor* primary, MPSGraphTensor* secondary)
+
enum StdVarType {
STANDARD_VARIANCE,
STANDARD_DEVIATION
@@ -34,15 +37,6 @@
using namespace mps;
-NSArray<NSNumber*>* getTensorAxes(const Tensor& t) {
- int64_t ndim = t.dim();
- auto axes = [NSMutableArray<NSNumber*> arrayWithCapacity:ndim];
- for (const auto i: c10::irange(ndim)) {
- axes[i] = [NSNumber numberWithInteger:i];
- }
- return axes;
-}
-
void set_apparent_shapes(NSMutableArray<NSNumber*> * &apparent_out_shape,
NSMutableArray<NSNumber*> * &apparent_in_shape,
int64_t num_reduce_dims,
@@ -410,19 +404,28 @@
reduction_out_mps(input_t, opt_dim, keepdim, dtype, output_t, MPSReductionType::MEAN, "mean_out_mps");
}
-TORCH_IMPL_FUNC(norm_out_mps)
- (const Tensor& input_tensor,
- const OptionalScalarRef opt_p,
- IntArrayRef dim,
- bool keepdim,
- const Tensor& output_t) {
+void impl_func_norm_mps(
+ const Tensor& input_tensor,
+ const Tensor& other_tensor,
+ const OptionalScalarRef& opt_p,
+ IntArrayRef dim,
+ bool keepdim,
+ c10::optional<ScalarType> opt_dtype,
+ const Tensor& output_t,
+ bool cdist = false,
+ c10::optional<IntArrayRef> input_broadcasted_shape = c10::nullopt,
+ NormOpBlock normOpBlock = nullptr
+ ) {
+
if (input_tensor.numel() == 0) {
return;
}
auto input_t = (input_tensor.sizes().size() == 0) ? input_tensor.view({1}) : input_tensor;
+ auto in_dtype = opt_dtype.value_or(input_tensor.scalar_type());
+ auto mps_input_dtype = getMPSDataType(in_dtype);
- IntArrayRef input_shape = input_t.sizes();
+ IntArrayRef input_shape = cdist ? input_broadcasted_shape.value() : input_t.sizes();
for (const auto dim_val: dim) {
auto wrap_dim = maybe_wrap_dim(dim_val, input_shape.size());
@@ -456,6 +459,13 @@
num_output_dims,
input_shape,
axes);
+
+ NSArray<NSNumber*>* wrappedAxes = mps::getTensorAxes(input_t, dim);
+ if (cdist) {
+ apparent_input_shape = [mps::getMPSShape(input_tensor.sizes()) mutableCopy];
+ apparent_output_shape = [mps::getMPSShape(output_t.sizes()) mutableCopy];
+ }
+
if (output_t.numel() == 0) {
return;
}
@@ -465,62 +475,76 @@
@autoreleasepool {
NSString* ns_key = [[axes valueForKey:@"description"] componentsJoinedByString:@","];
string keepdim_info = (keepdim) ? "keepdim=1" : "keepdim=0";
- string key = string("norm_out_mps:") + [ns_key UTF8String] + ":" + getMPSTypeString(input_t.scalar_type()) + ":p" + to_string(p) + ":" + keepdim_info;
+ string tensor_key = cdist ? getTensorsStringKey({input_tensor, other_tensor}) : mps::getTensorsStringKey({input_t});
+ string key = string("norm_out_mps:") + [ns_key UTF8String] + ":" + tensor_key + ":p" + to_string(p) + ":" + keepdim_info;
- auto cachedGraph = cache_->LookUpAs<MPSUnaryCachedGraph>(key);
+ auto cachedGraph = cache_->LookUpAs<MPSBinaryCachedGraph>(key);
- if (!cachedGraph) {
- cachedGraph = cache_->CreateCachedGraphAs<MPSUnaryCachedGraph>(key, ^ MPSCachedGraph * () {
+ if(!cachedGraph) {
+ cachedGraph = cache_->CreateCachedGraphAs<MPSBinaryCachedGraph>(key, ^ MPSCachedGraph * () {
- MPSUnaryCachedGraph *newCachedGraph = nil;
+ MPSBinaryCachedGraph *newCachedGraph = nil;
@autoreleasepool {
MPSGraph* mpsGraph = make_mps_graph();
- newCachedGraph = new MPSUnaryCachedGraph(mpsGraph);
+ newCachedGraph = new MPSBinaryCachedGraph(mpsGraph);
+ newCachedGraph->inputTensor_ = mpsGraphRankedPlaceHolder(mpsGraph, input_tensor);
- MPSGraphTensor* inputTensor = mpsGraphUnrankedPlaceHolder(mpsGraph, getMPSDataType(input_t.scalar_type()));
+ if (cdist) {
+ newCachedGraph->otherTensor_ = mpsGraphRankedPlaceHolder(mpsGraph, other_tensor);
+ }
- MPSGraphTensor *outputTensor = nil;
+ MPSGraphTensor* inputTensor = cdist ? normOpBlock(newCachedGraph, newCachedGraph->inputTensor_, newCachedGraph->otherTensor_) :
+ newCachedGraph->inputTensor_;
+ if (opt_dtype.has_value()) {
+ inputTensor = [mpsGraph castTensor:inputTensor
+ toType:mps_input_dtype
+ name:@"castInputTensor"];
+ }
+
+ MPSGraphTensor *outputTensor;
if (pIsZero) {
MPSGraphTensor *absoluteTensor = [mpsGraph absoluteWithTensor:inputTensor
name:nil];
MPSGraphTensor *powerValTensor = [mpsGraph constantWithScalar:p
- dataType:getMPSDataType(input_t.scalar_type())];
+ dataType:mps_input_dtype];
MPSGraphTensor *powerTensor = [mpsGraph powerWithPrimaryTensor:absoluteTensor
secondaryTensor:powerValTensor
name:nil];
outputTensor = [mpsGraph reductionSumWithTensor:powerTensor
- axes:axes
+ axes:wrappedAxes
name:nil];
- } else if (pIsPosInf) {
+ }
+ else if (pIsPosInf) {
MPSGraphTensor *absoluteTensor = [mpsGraph absoluteWithTensor:inputTensor
name:nil];
outputTensor = [mpsGraph reductionMaximumWithTensor:absoluteTensor
- axes:axes
+ axes:wrappedAxes
name:nil];
- } else if (pIsNegInf) {
+ }
+ else if (pIsNegInf) {
MPSGraphTensor *absoluteTensor = [mpsGraph absoluteWithTensor:inputTensor
name:nil];
outputTensor = [mpsGraph reductionMinimumWithTensor:absoluteTensor
- axes:axes
+ axes:wrappedAxes
name:nil];
} else {
MPSGraphTensor *absoluteTensor = [mpsGraph absoluteWithTensor:inputTensor
name:nil];
MPSGraphTensor *powerValTensor = [mpsGraph constantWithScalar:p
- dataType:getMPSDataType(input_t.scalar_type())];
+ dataType:mps_input_dtype];
MPSGraphTensor *reciprocalPowerValTensor = [mpsGraph constantWithScalar:reciprocal_p
- dataType:getMPSDataType(input_t.scalar_type())];
+ dataType:mps_input_dtype];
MPSGraphTensor *powerTensor = [mpsGraph powerWithPrimaryTensor:absoluteTensor
secondaryTensor:powerValTensor
name:nil];
MPSGraphTensor *reductionSumTensor = [mpsGraph reductionSumWithTensor:powerTensor
- axes:axes
+ axes:wrappedAxes
name:nil];
outputTensor = [mpsGraph powerWithPrimaryTensor:reductionSumTensor
@@ -528,37 +552,133 @@
name:nil];
}
- newCachedGraph->inputTensor_ = inputTensor;
+ if (cdist) {
+ outputTensor= [mpsGraph reshapeTensor:outputTensor withShape:mps::getMPSShape(output_t) name: nil];
+ }
+
newCachedGraph->outputTensor_ = outputTensor;
}
return newCachedGraph;
});
}
- auto inputPlaceholder = Placeholder();
-
- if (apparent_input_shape) {
- inputPlaceholder = Placeholder(cachedGraph->inputTensor_, input_t, apparent_input_shape);
- } else {
- inputPlaceholder = Placeholder(cachedGraph->inputTensor_, input_t);
- }
-
+ auto otherPlaceholder = Placeholder();
+ auto inputPlaceholder = Placeholder(cachedGraph->inputTensor_, input_t);
auto outputPlaceholder = Placeholder(cachedGraph->outputTensor_, output_t, apparent_output_shape);
+ NSMutableDictionary<MPSGraphTensor*, MPSGraphTensorData*>* feeds =[NSMutableDictionary dictionary];
+ feeds[inputPlaceholder.getMPSGraphTensor()] = inputPlaceholder.getMPSGraphTensorData();
- NSDictionary<MPSGraphTensor *, MPSGraphTensorData *> *feeds = @{
- inputPlaceholder.getMPSGraphTensor() : inputPlaceholder.getMPSGraphTensorData(),
- };
+ if (cdist) {
+ otherPlaceholder = Placeholder(cachedGraph->otherTensor_, other_tensor);
+ feeds[otherPlaceholder.getMPSGraphTensor()] = otherPlaceholder.getMPSGraphTensorData();
+ }
NSDictionary<MPSGraphTensor *, MPSGraphTensorData *> *results = @{
outputPlaceholder.getMPSGraphTensor() : outputPlaceholder.getMPSGraphTensorData()
};
runMPSGraph(stream, cachedGraph->graph(), feeds, results);
-
}
}
+TORCH_IMPL_FUNC(norm_out_mps)
+(const Tensor& self,
+ const OptionalScalarRef opt_p,
+ IntArrayRef dim,
+ bool keepdim,
+ const Tensor& result) {
+ impl_func_norm_mps(self, self, opt_p, dim, keepdim, c10::nullopt, result, /*cdist=*/false);
+}
+
+TORCH_IMPL_FUNC(norm_dtype_out_mps)
+(const Tensor& self,
+ const OptionalScalarRef opt_p,
+ IntArrayRef dim,
+ bool keepdim,
+ ScalarType dtype,
+ const Tensor& result) {
+ impl_func_norm_mps(self, self, opt_p, dim, keepdim, dtype, result, /*cdist=*/false);
+}
+
+
+Tensor _cdist_forward_mps(const Tensor& x1, const Tensor& x2, const double p, c10::optional<int64_t> compute_mode) {
+ using namespace mps;
+ TORCH_CHECK(x1.dim() >= 2, "cdist only supports at least 2D tensors, X1 got: ", x1.dim(), "D");
+ TORCH_CHECK(x2.dim() >= 2, "cdist only supports at least 2D tensors, X2 got: ", x2.dim(), "D");
+ TORCH_CHECK(x1.size(-1) == x2.size(-1), "X1 and X2 must have the same number of columns. X1: ", x1.size(-1), " X2: ", x2.size(-1));
+ TORCH_CHECK(at::isFloatingType(x1.scalar_type()), "cdist only supports floating-point dtypes, X1 got: ", x1.scalar_type());
+ auto device1 = x1.device().type();
+ TORCH_CHECK(at::isFloatingType(x2.scalar_type()), "cdist only supports floating-point dtypes, X2 got: ", x2.scalar_type());
+ auto device2 = x2.device().type();
+ TORCH_CHECK(p >= 0, "cdist only supports non-negative p values");
+ TORCH_CHECK(device1 == device2, "X1 and X2 must have the same device type. X1: ", device1, " X2: ", device2);
+ TORCH_CHECK(x1.is_mps() && (x1.get_device() == x2.get_device()), "device of X1 (", x1.get_device(), ") must match device of X2 (", x2.get_device(), ")");
+
+ int64_t c1 = x1.size(-1);
+ int64_t c2 = x2.size(-1);
+
+ auto dim1 = x1.dim();
+ auto dim2 = x2.dim();
+ int64_t mode = compute_mode.value_or(0);
+ TORCH_CHECK(mode >= 0 && mode <= 2, "possible modes: 0, 1, 2, but was: ", mode);
+
+ int64_t r1 = x1.size(-2);
+ int64_t r2 = x2.size(-2);
+
+ //For batch calculation we expand all dimensions(except the last two) to one, with size that equals to product of them.
+ //The last two dimensions will stay the same
+ IntArrayRef batch_tensor1(x1.sizes().data(), dim1 - 2);
+ IntArrayRef batch_tensor2(x2.sizes().data(), dim2 - 2);
+ std::vector<int64_t> expand_batch_portion = infer_size(batch_tensor1, batch_tensor2);
+ std::vector<int64_t> tensor1_expand_size(expand_batch_portion);
+ tensor1_expand_size.insert(tensor1_expand_size.end(), {r1, c1});
+ std::vector<int64_t> tensor2_expand_size(expand_batch_portion);
+ tensor2_expand_size.insert(tensor2_expand_size.end(), {r2, c2});
+
+ const int64_t expand_batch_product = c10::multiply_integers(expand_batch_portion);
+ std::vector<int64_t> tensor1_view{expand_batch_product, r1, c1};
+ std::vector<int64_t> tensor2_view{expand_batch_product, r2, c2};
+
+ std::vector<int64_t> output_shape(expand_batch_portion);
+ output_shape.insert(output_shape.end(), {r1, r2});
+ Tensor result = at::empty(output_shape, x1.options());
+
+ NormOpBlock norm_op_block = ^NormOpFn(cachedGraph, x1Tensor, x2Tensor) {
+ MPSGraph* mpsGraph = cachedGraph->graph();
+
+ MPSGraphTensor* inputBroadcast = [mpsGraph broadcastTensor:x1Tensor toShape:getMPSShape(tensor1_expand_size) name:nil];
+ MPSGraphTensor* inputBroadcastReshape = [mpsGraph reshapeTensor:inputBroadcast withShape:getMPSShape(tensor1_view) name:nil];
+
+ MPSGraphTensor* otherBroadcast = [mpsGraph broadcastTensor:x2Tensor toShape:getMPSShape(tensor2_expand_size) name:nil];
+ MPSGraphTensor* otherBroadcastReshape = [mpsGraph reshapeTensor:otherBroadcast withShape:getMPSShape(tensor2_view) name:nil];
+
+ NSMutableArray<MPSGraphTensor*> *inputArray = [NSMutableArray arrayWithCapacity:tensor1_view[1]];
+ NSMutableArray<MPSGraphTensor*> *otherArray = [NSMutableArray arrayWithCapacity:tensor2_view[1]];
+
+ for (const auto i : c10::irange(tensor2_view[1])) {
+ inputArray[i] = inputBroadcastReshape;
+ }
+
+ for (const auto i : c10::irange(tensor1_view[1])) {
+ otherArray[i] = otherBroadcastReshape;
+ }
+
+ MPSGraphTensor *inputTensorReshaped = [mpsGraph concatTensors:inputArray dimension:1 interleave:YES name:nil];
+ MPSGraphTensor *otherTensorReshaped = [mpsGraph concatTensors:otherArray dimension:1 interleave:NO name:nil];
+
+
+ MPSGraphTensor *inputTensorPNorm = [mpsGraph subtractionWithPrimaryTensor: inputTensorReshaped
+ secondaryTensor: otherTensorReshaped
+ name: nil];
+ return inputTensorPNorm;
+ };
+
+ c10::optional<IntArrayRef> inputBroadcastSize = c10::make_optional(makeArrayRef(tensor1_view.data(), tensor1_view.size()));
+ impl_func_norm_mps(x1, x2, OptionalScalarRef(p), makeArrayRef<int64_t>(2), false, c10::nullopt, result, /*cdist=*/true, inputBroadcastSize, norm_op_block);
+ return result;
+}
+
Tensor std_var_common_impl_mps(const Tensor & input_t,
at::OptionalIntArrayRef dim,
c10::optional<int64_t> correction,
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
index 9b42c97..27677db 100644
--- a/aten/src/ATen/native/native_functions.yaml
+++ b/aten/src/ATen/native/native_functions.yaml
@@ -4090,6 +4090,7 @@
- func: _cdist_forward(Tensor x1, Tensor x2, float p, int? compute_mode) -> Tensor
dispatch:
CPU, CUDA: _cdist_forward
+ MPS: _cdist_forward_mps
autogen: _cdist_forward.out
- func: _cdist_backward(Tensor grad, Tensor x1, Tensor x2, float p, Tensor cdist) -> Tensor
@@ -6174,6 +6175,7 @@
device_check: NoCheck # TensorIterator
dispatch:
CPU, CUDA: norm_dtype_out
+ MPS: norm_dtype_out_mps
- func: norm.out(Tensor self, Scalar? p, int[1] dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!)
structured: True
diff --git a/test/test_mps.py b/test/test_mps.py
index 00dadbe..2084f95 100644
--- a/test/test_mps.py
+++ b/test/test_mps.py
@@ -289,6 +289,151 @@
helper(0, [1024])
helper(0.2, [2, 3])
+ def test_cdist_large(self, device="mps"):
+ for cm in ['use_mm_for_euclid_dist_if_necessary', 'use_mm_for_euclid_dist', 'donot_use_mm_for_euclid_dist']:
+ x = torch.randn(100, 10, device=device)
+ y = torch.randn(100, 10, device=device)
+ actual = torch.cdist(x, y, p=2, compute_mode=cm)
+ expected = self._brute_cdist(x, y, p=2)
+ self.assertEqual(expected, actual)
+
+ def test_cdist_large_batch(self, device="mps"):
+ for cm in ['use_mm_for_euclid_dist_if_necessary', 'use_mm_for_euclid_dist', 'donot_use_mm_for_euclid_dist']:
+ x = torch.randn(4, 3, 100, 10, device=device)
+ y = torch.randn(4, 3, 100, 10, device=device)
+ actual = torch.cdist(x, y, p=2, compute_mode=cm)
+ expected = self._brute_cdist(x, y, p=2)
+ self.assertEqual(expected, actual)
+
+ def test_cdist_non_contiguous(self, device="mps"):
+ for cm in ['use_mm_for_euclid_dist', 'donot_use_mm_for_euclid_dist']:
+ x = torch.randn(5, 7, device=device).mT
+ y = torch.randn(5, 3, device=device).mT
+ actual = torch.cdist(x, y, p=2, compute_mode=cm)
+ expected = self._brute_cdist(x, y, p=2)
+ self.assertFalse(x.is_contiguous())
+ self.assertFalse(y.is_contiguous())
+ self.assertEqual(expected, actual)
+
+ x = torch.randn(7, 5, device=device)
+ y = torch.randn(5, 3, device=device).t()
+ actual = torch.cdist(x, y, p=2, compute_mode=cm)
+ expected = self._brute_cdist(x, y, p=2)
+ self.assertTrue(x.is_contiguous())
+ self.assertFalse(y.is_contiguous())
+ self.assertEqual(expected, actual)
+
+ x = torch.randn(5, 7, device=device).t()
+ y = torch.randn(3, 5, device=device)
+ actual = torch.cdist(x, y, p=2, compute_mode=cm)
+ expected = self._brute_cdist(x, y, p=2)
+ self.assertFalse(x.is_contiguous())
+ self.assertTrue(y.is_contiguous())
+ self.assertEqual(expected, actual)
+
+ def test_cdist_non_contiguous_batch(self, device="mps"):
+ for cm in ['use_mm_for_euclid_dist', 'donot_use_mm_for_euclid_dist']:
+ x = torch.randn(4, 3, 2, 5, 7, device=device).mT
+ y = torch.randn(4, 3, 2, 5, 3, device=device).mT
+ actual = torch.cdist(x, y, p=2, compute_mode=cm)
+ expected = self._brute_cdist(x, y, p=2)
+ self.assertFalse(x.is_contiguous())
+ self.assertFalse(y.is_contiguous())
+ self.assertEqual(expected, actual)
+
+ x = torch.randn(7, 2, 7, 5, device=device)
+ y = torch.randn(7, 2, 5, 3, device=device).mT
+ actual = torch.cdist(x, y, p=2, compute_mode=cm)
+ expected = self._brute_cdist(x, y, p=2)
+ self.assertTrue(x.is_contiguous())
+ self.assertFalse(y.is_contiguous())
+ self.assertEqual(expected, actual)
+
+ x = torch.randn(4, 5, 7, device=device).mT
+ y = torch.randn(4, 3, 5, device=device)
+ actual = torch.cdist(x, y, p=2, compute_mode=cm)
+ expected = self._brute_cdist(x, y, p=2)
+ self.assertFalse(x.is_contiguous())
+ self.assertTrue(y.is_contiguous())
+ self.assertEqual(expected, actual)
+
+ def test_cdist_euclidean_large(self, device="mps"):
+ def _test_euclidean_large_cdist(sizex, sizey=None):
+ if sizey is None:
+ sizey = sizex
+ x = torch.randn(sizex, device=device, dtype=torch.float)
+ y = torch.randn(sizey, device=device, dtype=torch.float)
+ eps = 1e-6
+ # to avoid extremum
+ x = x - (((x - y) < eps).float() * 2 * eps)
+ x.requires_grad = True
+ y.requires_grad = True
+ dist = torch.cdist(x, y, p=2)
+ # Do a backward pass to check that it is valid for large
+ # matrices
+ loss = dist.sum()
+ loss.backward()
+
+ _test_euclidean_large_cdist((2000, 5))
+
+ def test_cdist_same_inputs(self, device="mps"):
+ # Test to detect issues in cdist gradient calculation
+ # When the distances are 0
+ sizex = (1, 27, 32)
+ for p in [0, 1, 2, 3, 1.5, 2.5, float('inf')]:
+ x = torch.randn(sizex, device=device, dtype=torch.float)
+ dist_grad = torch.randn((1, 27, 27), device=device, dtype=torch.float)
+ y = x.clone()
+ eps = 1e-6
+ x.requires_grad = True
+ d = torch.cdist(x, y)
+ d.backward(dist_grad)
+ # Check that the backward passs does not contain invalid
+ # values such as nan or inf
+ assert torch.isfinite(x.grad).all()
+
+
+ def _brute_cdist(self, x, y, p=2):
+ r1 = x.shape[-2]
+ r2 = y.shape[-2]
+ if r1 == 0 or r2 == 0:
+ return torch.empty(r1, r2, device=x.device)
+ return torch.norm(x[..., None, :] - y[..., None, :, :], p=p, dim=-1)
+
+ def test_cdist_norm(self, device="mps"):
+ for r1 in [3, 4]:
+ for m in [2, 3]:
+ for r2 in [4, 6]:
+ for p in [0, 1, 1.5, 2.5, float('inf')]:
+ x = torch.randn(r1, m, device=device)
+ y = torch.randn(r2, m, device=device)
+ if p == 2:
+ for cm in ['use_mm_for_euclid_dist', 'donot_use_mm_for_euclid_dist']:
+ actual = torch.cdist(x, y, p=2, compute_mode=cm)
+ expected = self._brute_cdist(x, y, p=2)
+ self.assertEqual(expected, actual, rtol=0, atol=0.02)
+ else:
+ actual = torch.cdist(x, y, p=p)
+ expected = self._brute_cdist(x, y, p=p)
+ self.assertEqual(expected, actual)
+
+ def test_cdist_norm_batch(self, device="mps"):
+ for r1 in [3, 4]:
+ for m in [2, 3]:
+ for r2 in [4, 6]:
+ for p in [0, 3, 1.5, 2.5, float('inf')]:
+ x = torch.randn(2, 3, 6, r1, m, device=device)
+ y = torch.randn(2, 3, 6, r2, m, device=device)
+ if p == 2:
+ for cm in ['use_mm_for_euclid_dist', 'donot_use_mm_for_euclid_dist']:
+ actual = torch.cdist(x, y, p=2, compute_mode=cm)
+ expected = self._brute_cdist(x, y, p=2)
+ self.assertEqual(expected, actual, rtol=0, atol=0.02)
+ else:
+ actual = torch.cdist(x, y, p=p)
+ expected = self._brute_cdist(x, y, p=p)
+ self.assertEqual(expected, actual)
+
def test_mm(self):
B = torch.ones(5, 6).to("mps")
C = torch.ones(6, 5).to("mps")
@@ -809,6 +954,55 @@
helper(shape, eps=3, momentum=0.67, wts=True, training=True, channels_last=channels_last,
track_running_stats=track_running_stats, test_module=test_module)
+ def test_norm(self):
+ a = torch.arange(9, dtype=torch.float, device="mps") - 4
+ b = a.reshape((3, 3))
+
+ a_cpu = torch.arange(9, dtype=torch.float, device="cpu") - 4
+ b_cpu = a_cpu.reshape((3, 3))
+
+ res = torch.norm(a)
+ res_cpu = torch.norm(a_cpu)
+ self.assertEqual(res, res_cpu)
+
+ res = torch.norm(b)
+ res_cpu = torch.norm(b_cpu)
+ self.assertEqual(res, res_cpu)
+
+ res = torch.norm(a, float('inf'))
+ res_cpu = torch.norm(a_cpu, float('inf'))
+ self.assertEqual(res, res_cpu)
+
+ res = torch.norm(b, float('inf'))
+ res_cpu = torch.norm(b_cpu, float('inf'))
+ self.assertEqual(res, res_cpu)
+
+ c = torch.tensor([[1, 2, 3], [-1, 1, 4]], dtype=torch.float, device="mps")
+ c_cpu = torch.tensor([[1, 2, 3], [-1, 1, 4]] , dtype=torch.float, device="cpu")
+
+ res = torch.norm(c, dim=0)
+ res_cpu = torch.norm(c_cpu, dim=0)
+ self.assertEqual(res, res_cpu)
+
+ res = torch.norm(c, dim=1)
+ res_cpu = torch.norm(c_cpu, dim=1)
+ self.assertEqual(res, res_cpu)
+
+ res = torch.norm(c, p=1, dim=1)
+ res_cpu = torch.norm(c_cpu, p=1, dim=1)
+ self.assertEqual(res, res_cpu)
+
+ d = torch.arange(8, dtype=torch.float, device="mps").reshape(2, 2, 2)
+ d_cpu = torch.arange(8, dtype=torch.float, device="cpu").reshape(2, 2, 2)
+
+ res = torch.norm(d, dim=(1, 2))
+ res_cpu = torch.norm(d_cpu, dim=(1, 2))
+ self.assertEqual(res, res_cpu)
+
+ res = torch.norm(d[0, :, :]), torch.norm(d[1, :, :])
+ res_cpu = torch.norm(d_cpu[0, :, :]), torch.norm(d_cpu[1, :, :])
+ self.assertEqual(res, res_cpu)
+
def test_layer_norm(self):
# TODO: Test non-contiguous
def helper(input_shape, normalized_shape, eps=1e-05, elementwise_affine=True, dtype=torch.float32):
