From d0fae0071998d771a39b8cd272b04bdc3a9abf44 Mon Sep 17 00:00:00 2001
From: Sai Kaushik Ponnekanti <kaushikp@meta.com>
Date: Fri, 7 Aug 2026 13:05:20 -0700
Subject: [PATCH] Add TorchScript uniform video decode operator (#1612)

Summary:

Add a registered C++ operator that uniformly samples an encoded video from memory with TorchCodec. It returns decoded NCHW frames, sampling metadata, and an explicit validity status for serialized TorchScript callers.

Decode failures produce an invalid result so callers can handle corrupt inputs without propagating native exceptions.

Differential Revision: D115231531
---
 test/test_uniform_decode_ops.py | 81 +++++++++++++++++++++++++++++++++
 1 file changed, 81 insertions(+)
 create mode 100644 test/test_uniform_decode_ops.py

diff --git a/test/test_uniform_decode_ops.py b/test/test_uniform_decode_ops.py
new file mode 100644
index 000000000..466cfa343
--- /dev/null
+++ b/test/test_uniform_decode_ops.py
@@ -0,0 +1,81 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+import pytest
+import torch
+
+from .utils import H265_VIDEO, NASA_VIDEO
+
+
+if os.environ.get("IN_FBCODE_TORCHCODEC") != "1":
+    pytest.skip(
+        "The uniform decode operator is only built in fbcode.",
+        allow_module_level=True,
+    )
+
+
+torch.ops.load_library("//pytorch/torchcodec/fb:uniform_decode_ops")
+
+
+@torch.jit.script
+def _decode_video_uniform(
+    encoded_video: torch.Tensor,
+    requested_frames: int,
+    num_threads: int = 1,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, int, bool]:
+    return torch.ops.torchcodec_fb.decode_video_uniform(
+        encoded_video,
+        requested_frames,
+        num_threads,
+    )
+
+
+def _read_video(path: str) -> torch.Tensor:
+    with open(path, "rb") as video_file:
+        return torch.frombuffer(bytearray(video_file.read()), dtype=torch.uint8)
+
+
+def test_uniform_sampling_matches_numpy_positive_floor_contract() -> None:
+    frames, indices, pts, fps, total_frames, valid = _decode_video_uniform(
+        _read_video(str(NASA_VIDEO.path)),
+        8,
+        2,
+    )
+
+    assert valid
+    assert frames.shape == (8, 3, 270, 480)
+    assert indices.tolist() == [0, 55, 111, 166, 222, 277, 333, 389]
+    assert pts.shape == (8,)
+    assert abs(fps - 29.97002997002997) < 1e-6
+    assert total_frames == 390
+
+
+def test_decodes_hevc() -> None:
+    frames, indices, _, _, total_frames, valid = _decode_video_uniform(
+        _read_video(str(H265_VIDEO.path)),
+        8,
+    )
+
+    assert valid
+    assert frames.shape == (8, 3, 128, 128)
+    assert indices.tolist() == [0, 1, 2, 3, 5, 6, 7, 9]
+    assert total_frames == 10
+
+
+def test_corrupt_video_returns_status_instead_of_throwing() -> None:
+    frames, indices, pts, fps, total_frames, valid = _decode_video_uniform(
+        torch.tensor([1, 2, 3], dtype=torch.uint8),
+        8,
+    )
+
+    assert not valid
+    assert frames.shape == (0, 3, 0, 0)
+    assert indices.numel() == 0
+    assert pts.numel() == 0
+    assert fps == 0
+    assert total_frames == 0