| 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) |
|
|
