Startups.com·Machine Learning Engineer·Technical Phone Screen·Intermediate

Intermediate

Apr 2026

Summary

Interviewed for an ML Engineer role at Startups.com and got a pretty deep PyTorch implementation question. Not a vibe check, they wanted actual code and a real understanding of autograd internals.

Questions Asked (1)

Implement the ReLU activation function from scratch in PyTorch without using built-in ReLU utilities. Your solution should handle tensors of any shape, apply elementwise max(0, x), and support autograd with gradients of 1 where x > 0 and 0 elsewhere. Show both a plain tensor function version and a custom torch.autograd.Function with explicit forward and backward methods.

Technical Trade-offsAlgorithms & Data StructuresSystem Design

Author's notes

I knew ReLU conceptually but writing the autograd.Function version from memory was rougher than expected.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Start by clarifying the requirements: elementwise max(0, x), support for any shape, and autograd compatibility. Then present two implementations: a simple function using torch.clamp or torch.maximum, and a custom autograd.Function with explicit forward and backward. Explain how autograd works and why the custom function is useful for learning or customization.

Pro tip: Mention that while torch.clamp is concise, using torch.maximum(x, torch.zeros_like(x)) avoids potential issues with in-place operations and is more explicit. Also, highlight that the custom Function's backward must return a tuple with None for non-tensor inputs.

1. Clarify requirements and constraints

Restate the problem: implement ReLU without built-in ReLU, handle any shape, elementwise max(0, x), and support autograd. Confirm that using other torch ops like clamp or maximum is allowed.

2. Implement plain tensor function

Write a function that takes a tensor and returns torch.maximum(x, torch.zeros_like(x)) or torch.clamp(x, min=0). Explain that this leverages PyTorch's autograd automatically.

3. Implement custom autograd.Function

Define a class inheriting from torch.autograd.Function with static forward and backward methods. In forward, save the input or output for backward, and return the ReLU output. In backward, compute gradient as grad_output * (input > 0).

4. Test and validate

Show how to test both implementations with a sample tensor, including gradient computation via .backward() and comparing gradients to expected values. Mention handling of edge cases like zero input.

5. Discuss trade-offs and use cases

Compare the two approaches: simplicity vs. control, performance implications, and when a custom Function might be needed (e.g., for custom gradients or debugging).

Key Points to Mention

Elementwise max(0, x) can be implemented using torch.maximum or torch.clamp.
Autograd automatically handles gradients for built-in operations.
Custom autograd.Function requires forward and backward methods; backward returns gradients w.r.t. inputs.
In backward, gradient is 1 where input > 0, else 0; multiply by grad_output for chain rule.
Use torch.zeros_like(x) to create a zero tensor of the same shape and device.
Test with a tensor requiring grad and verify gradients using .backward() and .grad.

AI-generated suggestions, not part of the candidate's original notes. May be inaccurate — verify before relying on them.