chore: import upstream snapshot with attribution

tests/python/log_benchmark/__init__.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import dataclasses
+
+import numpy as np
+import numpy.typing as npt
+
+MAX_INT64 = 2**63 - 1
+MAX_INT32 = 2**31 - 1
+
+
+@dataclasses.dataclass
+class Point3DInput:
+    positions: npt.NDArray[np.float32]
+    colors: npt.NDArray[np.uint32]
+    radii: npt.NDArray[np.float32]
+    label: str = "some label"
+
+    @classmethod
+    def prepare(cls, seed: int, num_points: int) -> Point3DInput:
+        rng = np.random.default_rng(seed=seed)
+
+        return cls(
+            positions=rng.integers(0, MAX_INT64, (num_points, 3)).astype(dtype=np.float32),
+            colors=rng.integers(0, MAX_INT32, num_points, dtype=np.uint32),
+            radii=rng.integers(0, MAX_INT64, num_points).astype(dtype=np.float32),
+        )
+
+
+@dataclasses.dataclass
+class Transform3DInput:
+    """Input data for Transform3D benchmark with translation and mat3x3."""
+
+    translations: npt.NDArray[np.float32]  # Shape: (num_time_steps, num_entities, 3)
+    mat3x3s: npt.NDArray[np.float32]  # Shape: (num_time_steps, num_entities, 3, 3)
+    num_entities: int
+    num_time_steps: int
+
+    @classmethod
+    def prepare(cls, seed: int, num_entities: int, num_time_steps: int) -> Transform3DInput:
+        rng = np.random.default_rng(seed=seed)
+
+        # Generate translations in range [0, 10)
+        translations = rng.random((num_time_steps, num_entities, 3), dtype=np.float32) * 10.0
+
+        # Generate mat3x3 values in range [-1, 1)
+        mat3x3s = rng.random((num_time_steps, num_entities, 3, 3), dtype=np.float32) * 2.0 - 1.0
+
+        return cls(
+            translations=translations,
+            mat3x3s=mat3x3s,
+            num_entities=num_entities,
+            num_time_steps=num_time_steps,
+        )