tensorstudio 1.0.1__tar.gz → 1.1.0__tar.gz

{tensorstudio-1.0.1 → tensorstudio-1.1.0}/.github/workflows/publish.yml

@@ -4,9 +4,6 @@ on:
   workflow_dispatch:
   release:
     types: [published]
-  push:
-    tags:
-      - "v*"
 
 permissions:
   contents: read
@@ -97,6 +94,8 @@ jobs:
     needs: [windows_wheels, linux_wheels, macos_wheels, build_sdist]
     runs-on: ubuntu-latest
     environment: pypi
+    env:
+      HAS_PYPI_TOKEN: ${{ secrets.PYPI_TOKEN != '' }}
     steps:
       - uses: actions/download-artifact@v4
         with:
@@ -107,7 +106,14 @@ jobs:
         run: |
           python -m pip install -U twine
           python -m twine check dist/*
-      - name: Publish to PyPI
+      - name: Publish to PyPI with configured API token
+        if: env.HAS_PYPI_TOKEN == 'true'
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          password: ${{ secrets.PYPI_TOKEN }}
+          skip-existing: true
+      - name: Publish to PyPI with Trusted Publishing
+        if: env.HAS_PYPI_TOKEN != 'true'
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           skip-existing: true

{tensorstudio-1.0.1 → tensorstudio-1.1.0}/CHANGELOG.md

@@ -1,5 +1,29 @@
 # Changelog
 
+## Unreleased
+
+## 1.1.0 - 2026-07-07
+
+- Added native CPU NCHW `conv2d` with stride, padding, dilation, optional bias,
+  and reverse-mode gradients for input, weight, and bias.
+- Added `tensorstudio.conv2d`, `tensorstudio.ops.conv2d`,
+  `tensorstudio.nn.functional.conv2d`, and `tensorstudio.nn.Conv2d`.
+- Added convolution tests, finite-difference autograd coverage, and a
+  standalone convolution benchmark script.
+- Added convolution rows to the full benchmark report and updated docs to mark
+  convolution as partially implemented rather than missing.
+- Added native CPU NCHW `max_pool2d` and `avg_pool2d` operations with autograd,
+  Python top-level and functional APIs, `nn.MaxPool2d`, `nn.AvgPool2d`, tests,
+  docs, and benchmark coverage.
+- Added native single-axis `sum`, `mean`, `max`, and `min` reductions with
+  `keepdims` and reverse-mode gradients.
+- Added a typed contiguous 2D fast path for common axis `sum`/`mean`
+  reductions.
+- Added `nn.softmax`, `nn.log_softmax`, `nn.cross_entropy`, and
+  `nn.CrossEntropyLoss` for small multiclass classification workloads.
+- Added native `astype`, `concat`, and `stack` operations with Python wrappers,
+  Tensor methods for casting, and autograd support for floating tensors.
+
 ## 1.0.1
 
 - Added `test_all.py` for one-command local release checks covering lint,

{tensorstudio-1.0.1 → tensorstudio-1.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: tensorstudio
-Version: 1.0.1
+Version: 1.1.0
 Summary: TensorStudio is a compact C++ tensor and autograd engine with a Python API for learning, experimentation, and lightweight ML workloads.
 Keywords: tensor,autograd,machine-learning,cpp,pybind11
 Author: TensorStudio contributors
@@ -65,7 +65,7 @@ Description-Content-Type: text/markdown
 TensorStudio is a compact C++ tensor and autograd engine with a Python API for
 learning, experimentation, and lightweight ML workloads.
 
-TensorStudio `1.0.1` is a CPU-only stable API foundation. It is eager-only,
+TensorStudio `1.1.0` is a CPU-only stable API foundation. It is eager-only,
 intentionally small, and not a replacement for mature ML frameworks.
 
 ## Install
@@ -152,6 +152,8 @@ d = ts.linspace(0.0, 1.0, 5)
 
 print(a.shape, a.strides, a.device, a.is_contiguous)
 print((b.clamp(0.2, 0.8) + 1).mean().item())
+print(b.sum(axis=1).tolist())
+print(ts.concat([b, b], axis=0).shape, b.astype("float64").dtype)
 print(c.tolist(), d.tolist())
 print(ts.zeros_like(b).shape, ts.randn_like(b, seed=11).dtype)
 ```
@@ -235,12 +237,15 @@ python benchmarks/benchmark_report.py
 columns for NumPy, TensorFlow, PyTorch, and JAX when those libraries are
 available locally.
 
-On one Windows CPython 3.10 run for `1.0.1`, TensorStudio beat NumPy on 13
-small activation/reduction benchmark cases and lost on 76 NumPy-comparable
-cases. Against PyTorch CPU `2.12.1+cpu`, TensorStudio won 71 local cases and
-lost 23. The strongest local wins were small eager operations where framework
-dispatch overhead dominates; larger matrix multiplication, larger transcendental
-activations, and larger autograd workloads remain faster in PyTorch and NumPy.
+On one Windows CPython 3.10 development run reporting `1.1.0`, TensorStudio
+beat NumPy on 23 small operation benchmark cases and lost on 80
+NumPy-comparable cases. Against PyTorch CPU `2.12.1+cpu`, TensorStudio won 74
+local cases and lost 34. The strongest local wins were small eager operations,
+small contiguous axis reductions, and the simple NumPy convolution/pooling
+references where framework dispatch or Python loops dominate; larger matrix
+multiplication, PyTorch convolution and pooling, larger axis reductions, larger
+transcendental activations, and larger autograd workloads remain faster in
+PyTorch and NumPy.
 See `benchmarks/results.md` for the full table, platform details, and exact
 timings.
 
@@ -248,11 +253,15 @@ Snapshot from that local run:
 
 | operation | shape | TensorStudio | NumPy | PyTorch CPU | TS vs NumPy | TS vs PyTorch |
 |---|---:|---:|---:|---:|---:|---:|
-| `sigmoid` | `(32,)` | 0.0018 ms | 0.0036 ms | 0.0575 ms | 2.0414x | 32.2110x |
-| `mean` | `(32,)` | 0.0016 ms | 0.0069 ms | 0.0112 ms | 4.4051x | 7.1262x |
-| `chain_relu` | `(128,)` | 0.0086 ms | 0.0039 ms | 0.0567 ms | 0.4488x | 6.5977x |
-| `matmul` | `(256, 256)` | 2.4215 ms | 0.4027 ms | 0.0832 ms | 0.1663x | 0.0343x |
-| `elementwise_backward` | `(1024,)` | 2.3527 ms | n/a | 0.1887 ms | n/a | 0.0802x |
+| `sigmoid` | `(32,)` | 0.0017 ms | 0.0036 ms | 0.0580 ms | 2.1201x | 33.8536x |
+| `mean` | `(32,)` | 0.0018 ms | 0.0078 ms | 0.0127 ms | 4.2927x | 7.0047x |
+| `sum_axis1` | `(16, 16)` | 0.0021 ms | 0.0029 ms | 0.0068 ms | 1.3727x | 3.2193x |
+| `chain_relu` | `(128,)` | 0.0086 ms | 0.0039 ms | 0.0559 ms | 0.4478x | 6.5042x |
+| `matmul` | `(256, 256)` | 4.1154 ms | 0.3679 ms | 0.0931 ms | 0.0894x | 0.0226x |
+| `conv2d_3x3_padding1` | `(1, 1, 8, 8)` | 0.1788 ms | 1.2241 ms | 0.0131 ms | 6.8478x | 0.0731x |
+| `max_pool2d_2x2` | `(1, 1, 16, 16)` | 0.0146 ms | 0.2649 ms | 0.0062 ms | 18.1659x | 0.4242x |
+| `avg_pool2d_2x2` | `(1, 1, 16, 16)` | 0.0139 ms | 0.5486 ms | 0.0052 ms | 39.5974x | 0.3748x |
+| `elementwise_backward` | `(1024,)` | 2.3885 ms | n/a | 0.1947 ms | n/a | 0.0815x |
 
 Speedup is `competitor median / TensorStudio median`, so values above `1.0x`
 favor TensorStudio.
@@ -318,9 +327,12 @@ tokens or print secrets.
 - No CUDA or Metal backend yet.
 - No BLAS-backed matrix multiplication yet.
 - No graph compiler or distributed runtime.
-- No convolution layers yet.
+- Convolution and pooling support are currently limited to CPU NCHW
+  `conv2d`, `max_pool2d`, and `avg_pool2d` style workloads.
+- Reductions support all-element or single-axis reductions, not tuple-axis
+  reductions yet.
 - No sparse tensors or advanced indexing.
-- Limited dtype casting.
+- Dtype casting is basic and does not include a full promotion/casting policy.
 - Experimental performance; benchmarks are local references only.
 - Pickle serialization is for trusted TensorStudio objects only.
 
@@ -328,8 +340,8 @@ tokens or print secrets.
 
 - CUDA backend
 - Graph/JIT mode
-- Convolution ops
-- Dataset utilities
+- Broader convolution ops, adaptive/global pooling, and image-model examples
+- Richer dataset utilities
 - Model zoo examples
 - ONNX import/export
 - Improved memory allocator

{tensorstudio-1.0.1 → tensorstudio-1.1.0}/README.md

@@ -7,7 +7,7 @@
 TensorStudio is a compact C++ tensor and autograd engine with a Python API for
 learning, experimentation, and lightweight ML workloads.
 
-TensorStudio `1.0.1` is a CPU-only stable API foundation. It is eager-only,
+TensorStudio `1.1.0` is a CPU-only stable API foundation. It is eager-only,
 intentionally small, and not a replacement for mature ML frameworks.
 
 ## Install
@@ -94,6 +94,8 @@ d = ts.linspace(0.0, 1.0, 5)
 
 print(a.shape, a.strides, a.device, a.is_contiguous)
 print((b.clamp(0.2, 0.8) + 1).mean().item())
+print(b.sum(axis=1).tolist())
+print(ts.concat([b, b], axis=0).shape, b.astype("float64").dtype)
 print(c.tolist(), d.tolist())
 print(ts.zeros_like(b).shape, ts.randn_like(b, seed=11).dtype)
 ```
@@ -177,12 +179,15 @@ python benchmarks/benchmark_report.py
 columns for NumPy, TensorFlow, PyTorch, and JAX when those libraries are
 available locally.
 
-On one Windows CPython 3.10 run for `1.0.1`, TensorStudio beat NumPy on 13
-small activation/reduction benchmark cases and lost on 76 NumPy-comparable
-cases. Against PyTorch CPU `2.12.1+cpu`, TensorStudio won 71 local cases and
-lost 23. The strongest local wins were small eager operations where framework
-dispatch overhead dominates; larger matrix multiplication, larger transcendental
-activations, and larger autograd workloads remain faster in PyTorch and NumPy.
+On one Windows CPython 3.10 development run reporting `1.1.0`, TensorStudio
+beat NumPy on 23 small operation benchmark cases and lost on 80
+NumPy-comparable cases. Against PyTorch CPU `2.12.1+cpu`, TensorStudio won 74
+local cases and lost 34. The strongest local wins were small eager operations,
+small contiguous axis reductions, and the simple NumPy convolution/pooling
+references where framework dispatch or Python loops dominate; larger matrix
+multiplication, PyTorch convolution and pooling, larger axis reductions, larger
+transcendental activations, and larger autograd workloads remain faster in
+PyTorch and NumPy.
 See `benchmarks/results.md` for the full table, platform details, and exact
 timings.
 
@@ -190,11 +195,15 @@ Snapshot from that local run:
 
 | operation | shape | TensorStudio | NumPy | PyTorch CPU | TS vs NumPy | TS vs PyTorch |
 |---|---:|---:|---:|---:|---:|---:|
-| `sigmoid` | `(32,)` | 0.0018 ms | 0.0036 ms | 0.0575 ms | 2.0414x | 32.2110x |
-| `mean` | `(32,)` | 0.0016 ms | 0.0069 ms | 0.0112 ms | 4.4051x | 7.1262x |
-| `chain_relu` | `(128,)` | 0.0086 ms | 0.0039 ms | 0.0567 ms | 0.4488x | 6.5977x |
-| `matmul` | `(256, 256)` | 2.4215 ms | 0.4027 ms | 0.0832 ms | 0.1663x | 0.0343x |
-| `elementwise_backward` | `(1024,)` | 2.3527 ms | n/a | 0.1887 ms | n/a | 0.0802x |
+| `sigmoid` | `(32,)` | 0.0017 ms | 0.0036 ms | 0.0580 ms | 2.1201x | 33.8536x |
+| `mean` | `(32,)` | 0.0018 ms | 0.0078 ms | 0.0127 ms | 4.2927x | 7.0047x |
+| `sum_axis1` | `(16, 16)` | 0.0021 ms | 0.0029 ms | 0.0068 ms | 1.3727x | 3.2193x |
+| `chain_relu` | `(128,)` | 0.0086 ms | 0.0039 ms | 0.0559 ms | 0.4478x | 6.5042x |
+| `matmul` | `(256, 256)` | 4.1154 ms | 0.3679 ms | 0.0931 ms | 0.0894x | 0.0226x |
+| `conv2d_3x3_padding1` | `(1, 1, 8, 8)` | 0.1788 ms | 1.2241 ms | 0.0131 ms | 6.8478x | 0.0731x |
+| `max_pool2d_2x2` | `(1, 1, 16, 16)` | 0.0146 ms | 0.2649 ms | 0.0062 ms | 18.1659x | 0.4242x |
+| `avg_pool2d_2x2` | `(1, 1, 16, 16)` | 0.0139 ms | 0.5486 ms | 0.0052 ms | 39.5974x | 0.3748x |
+| `elementwise_backward` | `(1024,)` | 2.3885 ms | n/a | 0.1947 ms | n/a | 0.0815x |
 
 Speedup is `competitor median / TensorStudio median`, so values above `1.0x`
 favor TensorStudio.
@@ -260,9 +269,12 @@ tokens or print secrets.
 - No CUDA or Metal backend yet.
 - No BLAS-backed matrix multiplication yet.
 - No graph compiler or distributed runtime.
-- No convolution layers yet.
+- Convolution and pooling support are currently limited to CPU NCHW
+  `conv2d`, `max_pool2d`, and `avg_pool2d` style workloads.
+- Reductions support all-element or single-axis reductions, not tuple-axis
+  reductions yet.
 - No sparse tensors or advanced indexing.
-- Limited dtype casting.
+- Dtype casting is basic and does not include a full promotion/casting policy.
 - Experimental performance; benchmarks are local references only.
 - Pickle serialization is for trusted TensorStudio objects only.
 
@@ -270,8 +282,8 @@ tokens or print secrets.
 
 - CUDA backend
 - Graph/JIT mode
-- Convolution ops
-- Dataset utilities
+- Broader convolution ops, adaptive/global pooling, and image-model examples
+- Richer dataset utilities
 - Model zoo examples
 - ONNX import/export
 - Improved memory allocator

{tensorstudio-1.0.1 → tensorstudio-1.1.0}/benchmark_all.py

@@ -8,6 +8,8 @@ from benchmarks.benchmark_report import run_benchmarks
 ALL_SECTIONS = {
     "elementwise",
     "matmul",
+    "conv2d",
+    "pooling",
     "reductions",
     "activations",
     "autograd",

tensorstudio-1.1.0/benchmarks/bench_conv2d.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from benchmark_report import run_benchmarks
+
+if __name__ == "__main__":
+    print(run_benchmarks({"conv2d"}, Path(__file__).with_name("results_conv2d.md")))

tensorstudio-1.1.0/benchmarks/bench_pooling.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from benchmark_report import run_benchmarks
+
+if __name__ == "__main__":
+    print(run_benchmarks({"pooling"}, Path(__file__).with_name("results_pooling.md")))

{tensorstudio-1.0.1 → tensorstudio-1.1.0}/benchmarks/benchmark_report.py

@@ -24,6 +24,8 @@ Operation = Callable[[Any, Any | None], Any]
 
 VECTOR_SHAPES = [(1,), (8,), (32,), (128,), (1024,), (4096,), (16384,)]
 MATRIX_SHAPES = [(16, 16), (64, 64), (128, 128), (256, 256)]
+CONV2D_SHAPES = [(1, 1, 8, 8), (4, 3, 16, 16)]
+POOL2D_SHAPES = [(1, 1, 16, 16), (4, 3, 32, 32)]
 DTYPE = np.float32
 
 
@@ -88,6 +90,14 @@ def _iterations(shape: tuple[int, ...], category: str) -> int:
         if n <= 16384:
             return 5
         return 2
+    if category == "conv2d":
+        if n <= 256:
+            return 50
+        return 10
+    if category == "pooling":
+        if n <= 256:
+            return 200
+        return 50
     if category in {"autograd", "training_loop"}:
         if n <= 128:
             return 100
@@ -297,14 +307,25 @@ def _log_external(value: Any, module: str) -> Any:
     return np.log(value)
 
 
-def _reduction_op(name: str) -> Callable[[Library, np.ndarray, np.ndarray], Callable[[], Any]]:
+def _reduction_op(
+    name: str,
+    axis: int | None = None,
+) -> Callable[[Library, np.ndarray, np.ndarray], Callable[[], Any]]:
     def build(library: Library, left: np.ndarray, _: np.ndarray) -> Callable[[], Any]:
         x = library.factory(left)
 
         def run() -> Any:
             if library.name == "TensorStudio":
-                return x.sum() if name == "sum" else x.mean()
-            return _finalize(x.sum() if name == "sum" else x.mean())
+                return x.sum(axis=axis) if name == "sum" else x.mean(axis=axis)
+            module = x.__class__.__module__
+            if module.startswith("torch"):
+                return _finalize(x.sum(dim=axis) if name == "sum" else x.mean(dim=axis))
+            if module.startswith("tensorflow"):
+                tf = importlib.import_module("tensorflow")
+                return _finalize(
+                    tf.reduce_sum(x, axis=axis) if name == "sum" else tf.reduce_mean(x, axis=axis)
+                )
+            return _finalize(x.sum(axis=axis) if name == "sum" else x.mean(axis=axis))
 
         return run
 
@@ -321,6 +342,167 @@ def _matmul_build(library: Library, left: np.ndarray, right: np.ndarray) -> Call
     return run
 
 
+def _numpy_conv2d_nchw(
+    input: np.ndarray,
+    weight: np.ndarray,
+    padding: tuple[int, int] = (1, 1),
+) -> np.ndarray:
+    batch, in_channels, input_h, input_w = input.shape
+    out_channels, _, kernel_h, kernel_w = weight.shape
+    padding_h, padding_w = padding
+    out_h = input_h + 2 * padding_h - kernel_h + 1
+    out_w = input_w + 2 * padding_w - kernel_w + 1
+    output = np.zeros((batch, out_channels, out_h, out_w), dtype=input.dtype)
+    for n in range(batch):
+        for oc in range(out_channels):
+            for oh in range(out_h):
+                for ow in range(out_w):
+                    acc = 0.0
+                    for ic in range(in_channels):
+                        for kh in range(kernel_h):
+                            ih = oh - padding_h + kh
+                            if ih < 0 or ih >= input_h:
+                                continue
+                            for kw in range(kernel_w):
+                                iw = ow - padding_w + kw
+                                if iw < 0 or iw >= input_w:
+                                    continue
+                                acc += float(input[n, ic, ih, iw] * weight[oc, ic, kh, kw])
+                    output[n, oc, oh, ow] = acc
+    return output
+
+
+def _conv2d_build(library: Library, left: np.ndarray, right: np.ndarray) -> Callable[[], Any]:
+    if library.name == "TensorStudio":
+        x = ts.from_numpy(left)
+        w = ts.from_numpy(right)
+
+        def run_ts() -> Any:
+            return ts.conv2d(x, w, padding=1)
+
+        return run_ts
+
+    if library.name == "NumPy":
+        x_np = left
+        w_np = right
+
+        def run_np() -> Any:
+            return _numpy_conv2d_nchw(x_np, w_np)
+
+        return run_np
+
+    if library.name == "PyTorch CPU" and library.available:
+        torch = importlib.import_module("torch")
+        functional = importlib.import_module("torch.nn.functional")
+        x_torch = torch.from_numpy(left.copy())
+        w_torch = torch.from_numpy(right.copy())
+
+        def run_torch() -> Any:
+            return functional.conv2d(x_torch, w_torch, padding=1)
+
+        return run_torch
+
+    if library.name == "TensorFlow CPU eager" and library.available:
+        tf = importlib.import_module("tensorflow")
+        x_tf = tf.constant(np.transpose(left, (0, 2, 3, 1)))
+        w_tf = tf.constant(np.transpose(right, (2, 3, 1, 0)))
+
+        def run_tf() -> Any:
+            return tf.nn.conv2d(x_tf, w_tf, strides=1, padding="SAME").numpy()
+
+        return run_tf
+
+    if library.name == "JAX CPU dispatch" and library.available:
+        jax = importlib.import_module("jax")
+        jnp = importlib.import_module("jax.numpy")
+        x_jax = jnp.array(left)
+        w_jax = jnp.array(right)
+
+        def run_jax() -> Any:
+            return jax.lax.conv_general_dilated(
+                x_jax,
+                w_jax,
+                window_strides=(1, 1),
+                padding=((1, 1), (1, 1)),
+                dimension_numbers=("NCHW", "OIHW", "NCHW"),
+            ).block_until_ready()
+
+        return run_jax
+
+    raise NotImplementedError
+
+
+def _numpy_pool2d_nchw(input: np.ndarray, mode: str) -> np.ndarray:
+    batch, channels, input_h, input_w = input.shape
+    kernel_h = 2
+    kernel_w = 2
+    stride_h = 2
+    stride_w = 2
+    out_h = (input_h - kernel_h) // stride_h + 1
+    out_w = (input_w - kernel_w) // stride_w + 1
+    output = np.empty((batch, channels, out_h, out_w), dtype=input.dtype)
+    for n in range(batch):
+        for c in range(channels):
+            for oh in range(out_h):
+                for ow in range(out_w):
+                    window = input[
+                        n,
+                        c,
+                        oh * stride_h : oh * stride_h + kernel_h,
+                        ow * stride_w : ow * stride_w + kernel_w,
+                    ]
+                    output[n, c, oh, ow] = window.max() if mode == "max" else window.mean()
+    return output
+
+
+def _pool2d_build(mode: str) -> Callable[[Library, np.ndarray, np.ndarray], Callable[[], Any]]:
+    def build(library: Library, left: np.ndarray, _: np.ndarray) -> Callable[[], Any]:
+        if library.name == "TensorStudio":
+            x = ts.from_numpy(left)
+
+            def run_ts() -> Any:
+                if mode == "max":
+                    return ts.max_pool2d(x, 2)
+                return ts.avg_pool2d(x, 2)
+
+            return run_ts
+
+        if library.name == "NumPy":
+            x_np = left
+
+            def run_np() -> Any:
+                return _numpy_pool2d_nchw(x_np, mode)
+
+            return run_np
+
+        if library.name == "PyTorch CPU" and library.available:
+            torch = importlib.import_module("torch")
+            functional = importlib.import_module("torch.nn.functional")
+            x_torch = torch.from_numpy(left.copy())
+
+            def run_torch() -> Any:
+                if mode == "max":
+                    return functional.max_pool2d(x_torch, 2)
+                return functional.avg_pool2d(x_torch, 2, count_include_pad=False)
+
+            return run_torch
+
+        if library.name == "TensorFlow CPU eager" and library.available:
+            tf = importlib.import_module("tensorflow")
+            x_tf = tf.constant(np.transpose(left, (0, 2, 3, 1)))
+
+            def run_tf() -> Any:
+                if mode == "max":
+                    return tf.nn.max_pool2d(x_tf, ksize=2, strides=2, padding="VALID").numpy()
+                return tf.nn.avg_pool2d(x_tf, ksize=2, strides=2, padding="VALID").numpy()
+
+            return run_tf
+
+        raise NotImplementedError
+
+    return build
+
+
 def _chain_build(library: Library, left: np.ndarray, _: np.ndarray) -> Callable[[], Any]:
     x = library.factory(left)
 
@@ -497,9 +679,20 @@ def _build_cases(sections: set[str]) -> list[BenchmarkCase]:
         for op in ["sum", "mean"]:
             add_case("reductions", op, shape, _reduction_op(op))
 
+    for shape in MATRIX_SHAPES:
+        add_case("reductions", "sum_axis1", shape, _reduction_op("sum", axis=1))
+        add_case("reductions", "mean_axis0", shape, _reduction_op("mean", axis=0))
+
     for shape in MATRIX_SHAPES:
         add_case("matmul", "matmul", shape, _matmul_build)
 
+    for shape in CONV2D_SHAPES:
+        add_case("conv2d", "conv2d_3x3_padding1", shape, _conv2d_build)
+
+    for shape in POOL2D_SHAPES:
+        add_case("pooling", "max_pool2d_2x2", shape, _pool2d_build("max"))
+        add_case("pooling", "avg_pool2d_2x2", shape, _pool2d_build("avg"))
+
     for shape in [(1,), (128,), (1024,)]:
         add_case("autograd", "scalar_backward", shape, _autograd_scalar_build)
         add_case("autograd", "elementwise_backward", shape, _autograd_elementwise_build)
@@ -524,6 +717,10 @@ def _case_data(shape: tuple[int, ...], operation: str) -> tuple[np.ndarray, np.n
     rng = np.random.default_rng(123)
     left = rng.standard_normal(shape).astype(DTYPE)
     right = rng.standard_normal(shape).astype(DTYPE)
+    if operation == "conv2d_3x3_padding1":
+        _, in_channels, _, _ = shape
+        out_channels = 4
+        right = rng.standard_normal((out_channels, in_channels, 3, 3)).astype(DTYPE)
     if operation in {"div", "log"}:
         left = np.abs(left) + 0.25
         right = np.abs(right) + 0.25
@@ -741,7 +938,16 @@ def main() -> None:
     parser.add_argument(
         "--section",
         action="append",
-        choices=["elementwise", "matmul", "reductions", "activations", "autograd", "training_loop"],
+        choices=[
+            "elementwise",
+            "matmul",
+            "conv2d",
+            "pooling",
+            "reductions",
+            "activations",
+            "autograd",
+            "training_loop",
+        ],
         help="Benchmark section to run. Repeat for multiple sections. Defaults to all sections.",
     )
     parser.add_argument(
@@ -753,7 +959,16 @@ def main() -> None:
     args = parser.parse_args()
     sections = set(
         args.section
-        or ["elementwise", "matmul", "reductions", "activations", "autograd", "training_loop"]
+        or [
+            "elementwise",
+            "matmul",
+            "conv2d",
+            "pooling",
+            "reductions",
+            "activations",
+            "autograd",
+            "training_loop",
+        ]
     )
     report = run_benchmarks(sections, args.output)
     print(report)