From c209771dff97312f2b1733898495d4fc283c639e Mon Sep 17 00:00:00 2001 From: Junming Huang Date: Sun, 15 Mar 2026 11:46:38 -0700 Subject: [PATCH] Add Tanh & Soft-Capping exercise (#41) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add Problem #41 covering tanh forward, manual backward, and logit soft-capping (used in Gemma 2, Griffin). Includes numerically stable implementation using 2/(1+exp(-2x))-1 to handle large inputs. - templates/41_tanh.ipynb: template with 3 function stubs - solutions/41_tanh_solution.ipynb: working solution (stable formula) - torch_judge/tasks/tanh.py: 6 judge tests (accuracy, bounds, shape, gradient, manual backward, soft-capping) - README.md: add row to Fundamentals table, bump badge 40 โ†’ 41 Co-Authored-By: Claude Opus 4.6 --- README.md | 5 +- solutions/41_tanh_solution.ipynb | 83 ++++++++++++++++++ templates/41_tanh.ipynb | 144 +++++++++++++++++++++++++++++++ torch_judge/tasks/tanh.py | 77 +++++++++++++++++ 4 files changed, 307 insertions(+), 2 deletions(-) create mode 100644 solutions/41_tanh_solution.ipynb create mode 100644 templates/41_tanh.ipynb create mode 100644 torch_judge/tasks/tanh.py diff --git a/README.md b/README.md index e2b566b..af88402 100644 --- a/README.md +++ b/README.md @@ -27,7 +27,7 @@ Practice implementing operators and architectures from scratch โ€” the exact ski [![GitHub stars](https://img.shields.io/github/stars/duoan/TorchCode?style=social)](https://github.com/duoan/TorchCode) [![GitHub Container Registry](https://img.shields.io/badge/ghcr.io-TorchCode-blue?style=flat-square&logo=github)](https://ghcr.io/duoan/torchcode) [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Spaces-TorchCode-blue?style=flat-square)](https://huggingface.co/spaces/duoan/TorchCode) -![Problems](https://img.shields.io/badge/problems-40-orange?style=flat-square) +![Problems](https://img.shields.io/badge/problems-41-orange?style=flat-square) ![GPU](https://img.shields.io/badge/GPU-not%20required-brightgreen?style=flat-square) [![Star History Chart](https://api.star-history.com/svg?repos=duoan/TorchCode&type=Date)](https://star-history.com/#duoan/TorchCode&Date) @@ -44,7 +44,7 @@ TorchCode gives you a **structured practice environment** with: | | Feature | | |---|---|---| -| ๐Ÿงฉ | **40 curated problems** | The most frequently asked PyTorch interview topics | +| ๐Ÿงฉ | **41 curated problems** | The most frequently asked PyTorch interview topics | | โš–๏ธ | **Automated judge** | Correctness checks, gradient verification, and timing | | ๐ŸŽจ | **Instant feedback** | Colored pass/fail per test case, just like competitive programming | | ๐Ÿ’ก | **Hints when stuck** | Nudges without full spoilers | @@ -114,6 +114,7 @@ The bread and butter of ML coding interviews. You'll be asked to write these wit | 17 | Dropout Open In Colab | `MyDropout` (nn.Module) | ![Easy](https://img.shields.io/badge/Easy-4CAF50?style=flat-square) | ๐Ÿ”ฅ | Train/eval mode, inverted scaling | | 18 | Embedding Open In Colab | `MyEmbedding` (nn.Module) | ![Easy](https://img.shields.io/badge/Easy-4CAF50?style=flat-square) | ๐Ÿ”ฅ | Lookup table, `weight[indices]` | | 19 | GELU Open In Colab | `my_gelu(x)` | ![Easy](https://img.shields.io/badge/Easy-4CAF50?style=flat-square) | โญ | Gaussian error linear unit, `torch.erf` | +| 41 | Tanh Open In Colab | `my_tanh(x)`, `tanh_backward(...)`, `soft_cap_logits(...)` | ![Easy](https://img.shields.io/badge/Easy-4CAF50?style=flat-square) | โญ | Activation functions, backprop, logit soft-capping | | 20 | Kaiming Init Open In Colab | `kaiming_init(weight)` | ![Easy](https://img.shields.io/badge/Easy-4CAF50?style=flat-square) | โญ | `std = sqrt(2/fan_in)`, variance scaling | | 21 | Gradient Clipping Open In Colab | `clip_grad_norm(params, max_norm)` | ![Easy](https://img.shields.io/badge/Easy-4CAF50?style=flat-square) | โญ | Norm-based clipping, direction preservation | | 31 | Gradient Accumulation Open In Colab | `accumulated_step(model, opt, ...)` | ![Easy](https://img.shields.io/badge/Easy-4CAF50?style=flat-square) | ๐Ÿ’ก | Micro-batching, loss scaling | diff --git a/solutions/41_tanh_solution.ipynb b/solutions/41_tanh_solution.ipynb new file mode 100644 index 0000000..2c97e24 --- /dev/null +++ b/solutions/41_tanh_solution.ipynb @@ -0,0 +1,83 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/41_tanh_solution.ipynb)\n\n", + "# Solution: Tanh, Backward & Soft-Capping\n", + "\n", + "$$\\text{tanh}(x) = \\frac{e^x - e^{-x}}{e^x + e^{-x}}$$\n", + "\n", + "$$\\frac{d}{dx}\\text{tanh}(x) = 1 - \\text{tanh}^2(x)$$\n", + "\n", + "$$\\text{soft\\_cap}(x) = \\text{cap} \\cdot \\text{tanh}\\left(\\frac{x}{\\text{cap}}\\right)$$" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n", + "try:\n", + " import google.colab\n", + " get_ipython().run_line_magic('pip', 'install -q torch-judge')\n", + "except ImportError:\n", + " pass" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import torch" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "# โœ… SOLUTION\n\ndef my_tanh(x: torch.Tensor) -> torch.Tensor:\n # Equivalent to (e^x - e^-x)/(e^x + e^-x), but numerically stable\n # Divide numerator & denominator by e^x โ†’ 2ยทsigmoid(2x) - 1\n return 2.0 / (1.0 + torch.exp(-2.0 * x)) - 1.0\n\n\ndef tanh_backward(grad_output: torch.Tensor, tanh_output: torch.Tensor) -> torch.Tensor:\n return grad_output * (1 - tanh_output ** 2)\n\n\ndef soft_cap_logits(logits: torch.Tensor, cap: float) -> torch.Tensor:\n return cap * my_tanh(logits / cap)" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Verify\n", + "x = torch.tensor([-2., -1., 0., 1., 2.])\n", + "print('my_tanh:', my_tanh(x))\n", + "print('ref: ', torch.tanh(x))\n", + "\n", + "t = my_tanh(x)\n", + "print('backward:', tanh_backward(torch.ones_like(t), t))\n", + "\n", + "logits = torch.tensor([-50., 0., 50.])\n", + "print('soft_cap(cap=30):', soft_cap_logits(logits, 30.0))\n", + "\n", + "# Run judge\n", + "from torch_judge import check\n", + "check('tanh')" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/templates/41_tanh.ipynb b/templates/41_tanh.ipynb new file mode 100644 index 0000000..82356dc --- /dev/null +++ b/templates/41_tanh.ipynb @@ -0,0 +1,144 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/41_tanh.ipynb)\n\n", + "# ๐ŸŸข Easy: Tanh, Backward & Soft-Capping\n", + "\n", + "Implement the **tanh** activation function, its **backward pass**, and a **logit soft-capping** function.\n", + "\n", + "### Part 1 โ€” Forward\n", + "\n", + "$$\\text{tanh}(x) = \\frac{e^x - e^{-x}}{e^x + e^{-x}}$$\n", + "\n", + "```python\n", + "def my_tanh(x: torch.Tensor) -> torch.Tensor: ...\n", + "```\n", + "\n", + "### Part 2 โ€” Backward\n", + "\n", + "The derivative of tanh has an elegant property โ€” it can be expressed in terms of its own output:\n", + "\n", + "$$\\frac{d}{dx}\\text{tanh}(x) = 1 - \\text{tanh}^2(x)$$\n", + "\n", + "```python\n", + "def tanh_backward(grad_output: torch.Tensor, tanh_output: torch.Tensor) -> torch.Tensor: ...\n", + "```\n", + "\n", + "### Part 3 โ€” Soft-Capping (Gemma 2)\n", + "\n", + "Modern models like Gemma 2 use tanh to **soft-cap** logits, smoothly bounding them to $(-\\text{cap}, +\\text{cap})$:\n", + "\n", + "$$\\text{soft\\_cap}(x) = \\text{cap} \\cdot \\text{tanh}\\left(\\frac{x}{\\text{cap}}\\right)$$\n", + "\n", + "```python\n", + "def soft_cap_logits(logits: torch.Tensor, cap: float) -> torch.Tensor: ...\n", + "```\n", + "\n", + "### Rules\n", + "- Do **NOT** use `torch.tanh`, `F.tanh`, `torch.nn.Tanh`, or any built-in tanh\n", + "- Must support autograd (gradients should flow through `my_tanh`)\n", + "- `tanh_backward` should be a **manual** computation, not using autograd\n", + "- `soft_cap_logits` should use your `my_tanh`\n", + "\n", + "### Example\n", + "```\n", + "Input: tensor([-2., -1., 0., 1., 2.])\n", + "tanh: tensor([-0.9640, -0.7616, 0.0000, 0.7616, 0.9640])\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n", + "try:\n", + " import google.colab\n", + " get_ipython().run_line_magic('pip', 'install -q torch-judge')\n", + "except ImportError:\n", + " pass" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import torch" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# โœ๏ธ YOUR IMPLEMENTATION HERE\n", + "\n", + "def my_tanh(x: torch.Tensor) -> torch.Tensor:\n", + " \"\"\"Part 1: Implement tanh from scratch using exp.\"\"\"\n", + " pass # Replace this\n", + "\n", + "\n", + "def tanh_backward(grad_output: torch.Tensor, tanh_output: torch.Tensor) -> torch.Tensor:\n", + " \"\"\"Part 2: Manual backward โ€” compute gradient given upstream grad and tanh output.\"\"\"\n", + " pass # Replace this\n", + "\n", + "\n", + "def soft_cap_logits(logits: torch.Tensor, cap: float) -> torch.Tensor:\n", + " \"\"\"Part 3: Soft-cap logits to (-cap, +cap) using your my_tanh.\"\"\"\n", + " pass # Replace this" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# ๐Ÿงช Debug\n", + "x = torch.tensor([-2., -1., 0., 1., 2.])\n", + "print('my_tanh:', my_tanh(x))\n", + "print('ref: ', torch.tanh(x))\n", + "\n", + "# Test backward\n", + "t = my_tanh(x)\n", + "print('backward:', tanh_backward(torch.ones_like(t), t))\n", + "\n", + "# Test soft-capping\n", + "logits = torch.tensor([-50., 0., 50.])\n", + "print('soft_cap(cap=30):', soft_cap_logits(logits, 30.0))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# โœ… SUBMIT\n", + "from torch_judge import check\n", + "check('tanh')" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/torch_judge/tasks/tanh.py b/torch_judge/tasks/tanh.py new file mode 100644 index 0000000..506ecba --- /dev/null +++ b/torch_judge/tasks/tanh.py @@ -0,0 +1,77 @@ +"""Tanh activation, backward, and soft-capping task.""" + +TASK = { + "title": "Tanh, Backward & Soft-Capping", + "difficulty": "Easy", + "function_name": "my_tanh", + "hint": "tanh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x)). The derivative is 1 - tanh(x)^2 โ€” note it depends on the *output*, not the input. For soft-capping: cap * tanh(logits / cap).", + "tests": [ + { + "name": "Matches torch.tanh", + "code": """ +import torch +torch.manual_seed(0) +x = torch.randn(4, 8) +out = {fn}(x) +ref = torch.tanh(x) +assert torch.allclose(out, ref, atol=1e-5), f'Does not match torch.tanh' +""", + }, + { + "name": "tanh(0) = 0 and bounded output", + "code": """ +import torch +out_zero = {fn}(torch.tensor([0.0])) +assert torch.allclose(out_zero, torch.tensor([0.0]), atol=1e-7), f'tanh(0) should be 0, got {out_zero.item()}' +x_large = torch.tensor([100., -100.]) +out_large = {fn}(x_large) +assert (out_large.abs() <= 1.0 + 1e-5).all(), f'Output should be bounded in (-1, 1), got {out_large}' +""", + }, + { + "name": "Shape preservation", + "code": """ +import torch +x = torch.randn(2, 3, 4) +assert {fn}(x).shape == x.shape, 'Shape mismatch' +""", + }, + { + "name": "Gradient flow", + "code": """ +import torch +x = torch.randn(4, 8, requires_grad=True) +{fn}(x).sum().backward() +assert x.grad is not None and x.grad.shape == x.shape, 'Gradient issue' +""", + }, + { + "name": "Manual backward (tanh_backward)", + "code": """ +import torch +x = torch.randn(4, 8, requires_grad=True) +out = {fn}(x) +out.sum().backward() +autograd_grad = x.grad.clone() + +tanh_out = {fn}(x.detach()) +grad_output = torch.ones_like(tanh_out) +manual_grad = tanh_backward(grad_output, tanh_out) +assert torch.allclose(manual_grad, autograd_grad, atol=1e-5), f'tanh_backward does not match autograd' +""", + }, + { + "name": "Soft-capping bounds logits", + "code": """ +import torch +logits = torch.tensor([-50., -10., 0., 10., 50.]) +cap = 30.0 +capped = soft_cap_logits(logits, cap) +assert (capped.abs() < cap).all(), f'Soft-capped output should be within (-{cap}, {cap}), got {capped}' +assert torch.allclose(capped[2], torch.tensor(0.0), atol=1e-7), f'soft_cap(0) should be 0, got {capped[2]}' +ref = cap * torch.tanh(logits / cap) +assert torch.allclose(capped, ref, atol=1e-5), f'Does not match cap * tanh(logits / cap)' +""", + }, + ], +}