📘

Implement Custom ReLU Autograd Function

Problem Statement

Implement a custom ReLU activation function using torch.autograd.Function.

Background

PyTorch allows creating custom differentiable operations by subclassing torch.autograd.Function. You must implement:

forward(ctx, input): Compute the output and save tensors needed for backward
backward(ctx, grad_output): Compute gradients w.r.t. inputs

Your Task

The starter code defines a CustomReLU class and test harness. Implement the forward method (compute ReLU and save what's needed for backpropagation) and the backward method (compute the gradient of ReLU with respect to its input).

Output Format

The function returns a dictionary with "output" (ReLU values) and "grad" (input gradients).

Example:

Input:

None

Output:

{'output': [0.0, 0.0, 0.0, 1.0, 2.0], 'grad': [0.0, 0.0, 0.0, 1.0, 1.0]}

Reasoning:

We define a custom ReLU function CustomReLU that inherits from torch.autograd.Function, implementing the forward and backward methods. In forward, the input is saved and the output is computed as $max(0, x)$ using input.clamp(min=0).
The input tensor [-2.0, -1.0, 0.0, 1.0, 2.0] with requires_grad=True is created and passed through the custom ReLU function, resulting in the output [0.0, 0.0, 0.0, 1.0, 2.0].
The .sum().backward() method is called on the output, which computes the gradients of the input tensor. In the backward method, the gradient of the input is computed as $grad\_output$ where the input is greater than 0, and 0 otherwise, resulting in the gradient [0.0, 0.0, 0.0, 1.0, 1.0].
The output values and gradients are returned as a dictionary with keys "output" and "grad", resulting in the final output {'output': [0.0, 0.0, 0.0, 1.0, 2.0], 'grad': [0.0, 0.0, 0.0, 1.0, 1.0]}.

Constraints:

Must subclass torch.autograd.Function
Must use ctx.save_for_backward
Use @staticmethod for forward and backward

Editor

Python 3.13.1

Test Results

0/0

Run code to see test results.