Autograd Function with Multiple Outputs

mediumPROCustom Autograd FunctionsPyTorch Advanced

Problem Statement

Create a custom autograd function that returns multiple outputs from the forward pass.

Background

Some operations produce multiple results. Your custom function must handle gradients for each output separately, combining them correctly in the backward pass.

Your Task

The starter code defines a SplitAndScale class and test harness. Implement a function that separates positive and negative values from the input, scaling each appropriately. In the backward pass, handle the gradients from both outputs and combine them into a single gradient for the input.

Output Format

The function returns a dictionary with "positive", "negative", and "grad" keys.

Example

Input

None

Output

{'positive': [0.0, 0.0, 0.0, 4.0, 8.0], 'negative': [3.0, 1.0, 0.0, 0.0, 0.0], 'grad': [-1.0, -1.0, 0.0, 2.0, 2.0]}

Reasoning

The input x = [-3.0, -1.0, 0.0, 2.0, 4.0] is passed through the custom autograd function SplitAndScale.
The function splits the input into two tensors: positive = x.clamp(min=0) * 2 and negative = x.clamp(max=0) * -1, resulting in positive = [0.0, 0.0, 0.0, 4.0, 8.0] and negative = [3.0, 1.0, 0.0, 0.0, 0.0].
The function then computes the sum of both outputs, and calls backward to calculate the gradients. The gradients are computed as $grad = \begin{cases} grad_{pos} \cdot 2 & \text{if } x > 0 \\ grad_{neg} \cdot -1 & \text{if } x < 0 \\ 0 & \text{if } x = 0 \end{cases}$ , resulting in a gradient of [-1.0, -1.0, 0.0, 2.0, 2.0].
The final output is a dictionary containing the positive output, negative output, and the gradient of the input.

Constraints

Must return two tensors from forward
Handle gradients for both outputs in backward
Use ctx.save_for_backward

More from PyTorch Advanced

Implement SGD with Momentumhard Implement Custom SGD Optimizerhard Weight Decay for Specific Layersmedium Freeze Backbone, Train Headmedium Per-Layer Learning Rateseasy

Concept primerPREVIEW

Background Knowledge

The problem involves creating a custom autograd function in PyTorch, which is a fundamental component of the backpropagation algorithm used for training neural networks. Autograd functions are responsible for computing the forward pass and the backward pass of a neural network. The forward pass involves applying the function to the input data, while the backward pass involves computing the gradients of the loss with respect to the input data. In this problem, we need to define a custom autograd function that returns multiple outputs from the forward pass and handles gradients for each output separately.

To understand this problem, it's essential to have a solid grasp of PyTorch's autograd system, which is a dynamic computation graph that records all the operations performed on tensors. When a tensor is created with requires_grad=True, PyTorch automatically records all the operations performed on that tensor, allowing us to compute the gradients of the output with respect to the input using the chain rule. In this problem, we need to define a custom autograd function that takes into account the multiple outputs and computes the gradients correctly.

The problem also involves understanding tensor operations in PyTorch, such as clamp, which is used to limit the values of a tensor to a specific range. We also need to understand how to index and mask tensors to apply different operations to different parts of the tensor. Additionally, we need to understand how to combine gradients from multiple outputs to compute the final gradient of the input.

Algorithm/Approach

The general approach to solving this problem involves defining a custom autograd function that inherits from torch.autograd.Function. We need to override the forward method to define the forward pass of the function, which involves applying the clamp operation to the input tensor and returning two separate tensors. We also need to override the backward method to define the backward pass of the function, which involves computing the gradients of the loss with respect to the input tensor.

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test results