Are you even allowed to do these ops inplace?

[mayank31398]

Liger-Kernel/src/liger_kernel/ops/swiglu.py

Lines 59 to 60 in 58fd2bc

	tl.store(a_ptr + col_offsets, da_row, mask=mask)
	tl.store(b_ptr + col_offsets, db_row, mask=mask)

Lets take a custom autograd function:

class Exponential(torch.autograd.Function):
    def forward(ctx, x):
        out = torch.exp(x)
        ctx.save_for_backward(out)
        return out

    def backward(ctx, out_grad):
        out = ctx.saved_tensors
        x_grad = out_grad * out
        return x_grad

and if we have an op like swiglu that modifies inputs in backwards:

x = some tensor
x_exp = Exponential.apply(x)
y = swiglu(x_exp, x_exp)
loss = some_loss(y, target)

now during backprop, we would see incorrect behaviour right?
because the custom autograd function Exponential saves the output for backprop here instead of saving the input for backprop.

[mayank31398]

Hey, guys
Any clarification regarding this?