Is there a good way to modify some values in a pytorch tensor while preserving the autograd functionality?

Question

Sometimes I need to modify some of the values in a pytorch tensor. For example, given a tensor x, I need to multiply its positive part by 2 and multiply its negative part by 3:

import torch

x = torch.randn(1000, requires_grad=True)
x[x>0] = 2 * x[x>0]
x[x<0] = 3 * x[x<0]

y = x.sum()
y.backward()

However such inplace operations always break the graph for autograd:

Traceback (most recent call last):
  File "test_rep.py", line 4, in <module>
    x[x>0] = 2 * x[x>0]
RuntimeError: a leaf Variable that requires grad is being used in an in-place operation.

Therefore, so far I've been using the following workaround:

import torch

x = torch.randn(1000, requires_grad=True)
y = torch.zeros_like(x, device=x.device)

y[x>0] = 2 * x[x>0]
y[x<0] = 3 * x[x<0]

z = y.sum()
z.backward()

which results in manually creating new tensors. I wonder if there is a better way to do this.

Answer 1

How about like following?

import torch

x = torch.randn(1000, requires_grad=True)
x = torch.where(x>0, x*2, x)
x = torch.where(x<0, x*3, x)

y = x.sum()
y.backward()

Answer 2

There is a better way, at least for this specific case, based on how LeakyReLU works:

import torch

x = torch.randn(10, requires_grad=True)

y = 2 * torch.max(x, torch.tensor(0.0)) + 3 * torch.min(x, torch.tensor(0.0))

Some broadcasting will be used for 0-dim torch.tensor(0.0).

PyTorch does not work well with in-place operations, as the tensors are recorded on tape (and ops which created them). If you override some of those values (here tensors where x>0 or x<0), their history will be overridden instead of appended (as is the case with applying out of place operations on tensor, as above).

There is rarely ever a case to use inplace operations (unless you are really constrained by memory usage).