code: def calculate_shards(
      num_shards: int,
      tests: Sequence[TestRun],
      test_file_times: Dict[str, float],
      test_class_times: Optional[Dict[str, Dict[str, float]]],
      must_serial: Optional[Callable[[str], bool]] = None,
      sort_by_time: bool = True,
  ) -> List[Tuple[float, List[ShardedTest]]]:
      must_serial = must_serial or (lambda x: True)
  
      known_tests: Sequence[TestRun] = tests
      unknown_tests: Sequence[TestRun] = []
  
      if sort_by_time:
          known_tests = [
              x
              for x in tests
              if x.test_file in test_file_times
              or (test_class_times and x.test_file in test_class_times)
          ]
          unknown_tests = [x for x in tests if x not in known_tests]
  
      known_tests = get_with_pytest_shard(
          known_tests, test_file_times, test_class_times
      )
  
      if sort_by_time:
          known_tests = sorted(
              known_tests, key=lambda j: j.get_time(), reverse=True
          )
  
      sharded_jobs: List[ShardJob] = [
          ShardJob() for _ in range(num_shards)
      ]
      for test in known_tests:
          if must_serial(test.name):
              min_sharded_job = min(
                  sharded_jobs, key=lambda j: j.get_total_time()
              )
              min_sharded_job.serial.append(test)
          else:
              min_sharded_job = min(
                  sharded_jobs, key=lambda j: j.get_total_time()
              )
              min_sharded_job.parallel.append(test)
  
      # Round robin the unknown jobs starting with the smallest shard
      index = min(
          range(num_shards),
          key=lambda i: sharded_jobs[i].get_total_time(),
      )
      for unknown_test in unknown_tests:
          sharded_jobs[index].serial.append(
              ShardedTest(unknown_test, 1, 1, None)
          )
          index = (index + 1) % num_shards
      return [job.convert_to_tuple() for job in sharded_jobs]
  
file_path: pytorch/tools/testing/test_selections.py
imports: import math
  import os
  import subprocess
  from pathlib import Path
  from typing import Callable, Dict, List, Optional, Sequence, Set, Tuple
  from tools.stats.import_test_stats import get_disabled_tests, get_slow_tests
  from tools.testing.test_run import ShardedTest, TestRun
label_desc: Use image node_id 2 for calling a global function with example usage: calculate_shards(num_shards, tests, test_file_times, test_class_times, must_serial, sort_by_time) and returns: unknown