Add types to binary_trees and spectral_norm benchmarks

Summary: Continuing down the path of adding type annotations to our benchmark files.

Reviewed By: brittanyrey

Differential Revision: D96810282

fbshipit-source-id: 95ba94089e7ca523c89aa55bf7ffe29e55227ca3

Author: alexmalyshev

cinderx/benchmarks/binary_trees.py

@@ -1,5 +1,7 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 
+# pyre-strict
+
 """
 Binary trees benchmark.
 
@@ -12,32 +14,35 @@
 Benchmarks Game.
 """
 
+from __future__ import annotations
+
 import sys
 
 import cinderx.jit
 
 
-class TreeNode(object):
-    def __init__(self, left, right):
+class TreeNode:
+    def __init__(self, left: TreeNode | None, right: TreeNode | None) -> None:
         self.left = left
         self.right = right
 
 
-def make_tree(depth):
+def make_tree(depth: int) -> TreeNode:
     if depth <= 0:
         return TreeNode(None, None)
     depth -= 1
     return TreeNode(make_tree(depth), make_tree(depth))
 
 
-def check_tree(node):
+def check_tree(node: TreeNode) -> int:
     if node.left is None:
         return 1
+    # pyre-ignore[6]: Subtrees have to be non-None.
     return 1 + check_tree(node.left) + check_tree(node.right)
 
 
-class BinaryTrees(object):
-    def run(self, iterations):
+class BinaryTrees:
+    def run(self, iterations: int) -> bool:
         for _ in range(iterations):
             min_depth = 4
             max_depth = 17

cinderx/benchmarks/spectral_norm.py

@@ -1,5 +1,7 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 
+# pyre-strict
+
 """
 Spectral norm benchmark.
 
@@ -18,12 +20,12 @@
 import cinderx.jit
 
 
-def eval_A(i, j):
+def eval_A(i: int, j: int) -> float:
     ij = i + j
     return 1.0 / (ij * (ij + 1) // 2 + i + 1)
 
 
-def eval_A_times_u(u, n):
+def eval_A_times_u(u: list[float], n: int) -> list[float]:
     result = []
     for i in range(n):
         s = 0.0
@@ -33,7 +35,7 @@ def eval_A_times_u(u, n):
     return result
 
 
-def eval_At_times_u(u, n):
+def eval_At_times_u(u: list[float], n: int) -> list[float]:
     result = []
     for i in range(n):
         s = 0.0
@@ -43,11 +45,11 @@ def eval_At_times_u(u, n):
     return result
 
 
-def eval_AtA_times_u(u, n):
+def eval_AtA_times_u(u: list[float], n: int) -> list[float]:
     return eval_At_times_u(eval_A_times_u(u, n), n)
 
 
-def spectral_norm(n):
+def spectral_norm(n: int) -> float:
     u = [1.0] * n
     v = [0.0] * n
 
@@ -63,8 +65,8 @@ def spectral_norm(n):
     return (vBv / vv) ** 0.5
 
 
-class SpectralNorm(object):
-    def run(self, iterations):
+class SpectralNorm:
+    def run(self, iterations: int) -> bool:
         for _ in range(iterations):
             n = 1200
             result = spectral_norm(n)