|
|
import torch |
|
|
import torch.nn.functional as F |
|
|
|
|
|
from ._ops import add_op_namespace_prefix |
|
|
|
|
|
|
|
|
@torch.library.custom_op(add_op_namespace_prefix("silu_and_mul"), mutates_args=()) |
|
|
def _silu_and_mul(x: torch.Tensor) -> torch.Tensor: |
|
|
d = x.shape[-1] // 2 |
|
|
return F.silu(x[..., :d]) * x[..., d:] |
|
|
|
|
|
|
|
|
def backward(ctx, grad_output): |
|
|
x = ctx.saved_tensors[0] |
|
|
d = x.shape[-1] // 2 |
|
|
x1, x2 = x[..., :d], x[..., d:] |
|
|
sigmoid_x1 = torch.sigmoid(x1) |
|
|
silu_x1 = F.silu(x1) |
|
|
dsilu_dx1 = sigmoid_x1 + silu_x1 * (1 - sigmoid_x1) |
|
|
dx1 = grad_output * x2 * dsilu_dx1 |
|
|
dx2 = grad_output * silu_x1 |
|
|
return torch.cat([dx1, dx2], dim=-1) |
|
|
|
|
|
|
|
|
def setup_context(ctx, inputs, output): |
|
|
(x,) = inputs |
|
|
ctx.save_for_backward(x) |
|
|
|
|
|
|
|
|
_silu_and_mul.register_autograd(backward, setup_context=setup_context) |
|
|
|
|
|
|
|
|
@_silu_and_mul.register_fake |
|
|
def _(x: torch.Tensor) -> torch.Tensor: |
|
|
return x.new_empty(x.shape[0], x.shape[1] // 2) |
|
|
|
