Include scalar params in caching key (#94)

* Include scalar params in caching key * Add key for softplus backward; add test for scalar params
kulinseth · Oct 31, 2022 · 4169a72 · 4169a72
1 parent 590be21
commit 4169a72
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/aten/src/ATen/native/mps/operations/Activation.mm b/aten/src/ATen/native/mps/operations/Activation.mm
@@ -1474,7 +1474,7 @@ Tensor glu_backward_mps (const Tensor& grad_output,
  MPSScalar beta_scalar = getMPSScalar(beta, ScalarType::Float);;
 
  @autoreleasepool {
- string key = "softplus_out_mps:" + getTensorsStringKey({self});
+ string key = "softplus_out_mps:" + getTensorsStringKey({self}) + ":" + std::to_string(beta.to<double>()) + ":" + std::to_string(threshold.to<double>());
 
  CachedGraph* cachedGraph = static_cast<CachedGraph *>(cache_->LookUp(key));
  if(!cachedGraph) {
@@ -1576,7 +1576,7 @@ Tensor glu_backward_mps (const Tensor& grad_output,
  MPSStream* stream = getCurrentMPSStream();
 
  @autoreleasepool {
- string key = "softplus_backward_out_mps:" + getTensorsStringKey({grad_output, self});
+ string key = "softplus_backward_out_mps:" + getTensorsStringKey({grad_output, self}) + ":" + std::to_string(beta.to<double>()) + ":" + std::to_string(threshold.to<double>());
 
  CachedGraph* cachedGraph = static_cast<CachedGraph *>(cache_->LookUp(key));
  if(!cachedGraph) {

diff --git a/test/test_mps.py b/test/test_mps.py
@@ -3853,12 +3853,12 @@ def helper(shape, dim=0):
 
  # Test softplus
  def test_softplus(self):
- def helper(shape):
+ def helper(shape, beta=1, threshold=20):
  cpu_x = torch.randn(shape, device='cpu', dtype=torch.float, requires_grad=True)
  x = cpu_x.detach().clone().to('mps').requires_grad_()
 
- softplus_result = torch.nn.Softplus(beta=0.5, threshold=0.5)(x)
- softplus_result_cpu = torch.nn.Softplus(beta=0.5, threshold=0.5)(cpu_x)
+ softplus_result = torch.nn.Softplus(beta=beta, threshold=threshold)(x)
+ softplus_result_cpu = torch.nn.Softplus(beta=beta, threshold=threshold)(cpu_x)
 
  cpu_grad = torch.randn(softplus_result.shape)
  grad = cpu_grad.to('mps')
@@ -3871,7 +3871,9 @@ def helper(shape):
 
  # Test empty shape too
  for shape in [(), (2, 3), (10, 10), (2, 3, 4, 5)]:
- helper(shape)
+ for beta in [0.5, 1, 2, 3, 4]:
+ for threshold in [0.5, 20, 30, 40, 50]:
+ helper(shape, beta, threshold)
 
  # Test silu