content: import argparse
  import io
  import os
  import random
  import warnings
  import zipfile
  from abc import ABC, abstractmethod
  from contextlib import contextmanager
  from functools import partial
  from multiprocessing import cpu_count
  from multiprocessing.pool import ThreadPool
  from typing import Iterable, Optional, Tuple
  import yaml
  
  import numpy as np
  import requests
  import tensorflow.compat.v1 as tf
  from scipy import linalg
  from tqdm.auto import tqdm
  
  INCEPTION_V3_URL = "https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/classify_image_graph_def.pb"
  INCEPTION_V3_PATH = "classify_image_graph_def.pb"
  
  FID_POOL_NAME = "pool_3:0"
  FID_SPATIAL_NAME = "mixed_6/conv:0"
  
  REQUIREMENTS = f"This script has the following requirements: \n" \
                 'tensorflow-gpu>=2.0' + "\n" + 'scipy' + "\n" + "requests" + "\n" + "tqdm"
  
  
  def main():
      parser = argparse.ArgumentParser()
      parser.add_argument("--ref_batch", help="path to reference batch npz file")
      parser.add_argument("--sample_batch", help="path to sample batch npz file")
      args = parser.parse_args()
  
      config = tf.ConfigProto(
          allow_soft_placement=True  # allows DecodeJpeg to run on CPU in Inception graph
      )
      config.gpu_options.allow_growth = True
      evaluator = Evaluator(tf.Session(config=config))
  
      print("warming up TensorFlow...")
      # This will cause TF to print a bunch of verbose stuff now rather
      # than after the next print(), to help prevent confusion.
      evaluator.warmup()
  
      print("computing reference batch activations...")
      ref_acts = evaluator.read_activations(args.ref_batch)
      print("computing/reading reference batch statistics...")
      ref_stats, ref_stats_spatial = evaluator.read_statistics(args.ref_batch, ref_acts)
  
      print("computing sample batch activations...")
      sample_acts = evaluator.read_activations(args.sample_batch)
      print("computing/reading sample batch statistics...")
      sample_stats, sample_stats_spatial = evaluator.read_statistics(args.sample_batch, sample_acts)
  
      print("Computing evaluations...")
      is_ = evaluator.compute_inception_score(sample_acts[0])
      print("Inception Score:", is_)
      fid = sample_stats.frechet_distance(ref_stats)
      print("FID:", fid)
      sfid = sample_stats_spatial.frechet_distance(ref_stats_spatial)
      print("sFID:", sfid)
      prec, recall = evaluator.compute_prec_recall(ref_acts[0], sample_acts[0])
      print("Precision:", prec)
      print("Recall:", recall)
  
      savepath = '/'.join(args.sample_batch.split('/')[:-1])
      results_file = os.path.join(savepath,'evaluation_metrics.yaml')
      print(f'Saving evaluation results to "{results_file}"')
  
      results = {
          'IS': is_,
          'FID': fid,
          'sFID': sfid,
          'Precision:':prec,
          'Recall': recall
      }
  
      with open(results_file, 'w') as f:
          yaml.dump(results, f, default_flow_style=False)
  
  class InvalidFIDException(Exception):
      pass
  
  
  class FIDStatistics:
      def __init__(self, mu: np.ndarray, sigma: np.ndarray):
          self.mu = mu
          self.sigma = sigma
  
      def frechet_distance(self, other, eps=1e-6):
          """
          Compute the Frechet distance between two sets of statistics.
          """
          # https://github.com/bioinf-jku/TTUR/blob/73ab375cdf952a12686d9aa7978567771084da42/fid.py#L132
          mu1, sigma1 = self.mu, self.sigma
          mu2, sigma2 = other.mu, other.sigma
  
          mu1 = np.atleast_1d(mu1)
          mu2 = np.atleast_1d(mu2)
  
          sigma1 = np.atleast_2d(sigma1)
          sigma2 = np.atleast_2d(sigma2)
  
          assert (
              mu1.shape == mu2.shape
          ), f"Training and test mean vectors have different lengths: {mu1.shape}, {mu2.shape}"
          assert (
              sigma1.shape == sigma2.shape
          ), f"Training and test covariances have different dimensions: {sigma1.shape}, {sigma2.shape}"
  
          diff = mu1 - mu2
  
          # product might be almost singular
          covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
          if not np.isfinite(covmean).all():
              msg = (
                  "fid calculation produces singular product; adding %s to diagonal of cov estimates"
                  % eps
              )
              warnings.warn(msg)
              offset = np.eye(sigma1.shape[0]) * eps
              covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
  
          # numerical error might give slight imaginary component
          if np.iscomplexobj(covmean):
              if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
                  m = np.max(np.abs(covmean.imag))
                  raise ValueError("Imaginary component {}".format(m))
              covmean = covmean.real
  
          tr_covmean = np.trace(covmean)
  
          return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean
  
  
  class Evaluator:
      def __init__(
          self,
          session,
          batch_size=64,
          softmax_batch_size=512,
      ):
          self.sess = session
          self.batch_size = batch_size
          self.softmax_batch_size = softmax_batch_size
          self.manifold_estimator = ManifoldEstimator(session)
          with self.sess.graph.as_default():
              self.image_input = tf.placeholder(tf.float32, shape=[None, None, None, 3])
              self.softmax_input = tf.placeholder(tf.float32, shape=[None, 2048])
              self.pool_features, self.spatial_features = _create_feature_graph(self.image_input)
              self.softmax = _create_softmax_graph(self.softmax_input)
  
      def warmup(self):
          self.compute_activations(np.zeros([1, 8, 64, 64, 3]))
  
      def read_activations(self, npz_path: str) -> Tuple[np.ndarray, np.ndarray]:
          with open_npz_array(npz_path, "arr_0") as reader:
              return self.compute_activations(reader.read_batches(self.batch_size))
  
      def compute_activations(self, batches: Iterable[np.ndarray],silent=False) -> Tuple[np.ndarray, np.ndarray]:
          """
          Compute image features for downstream evals.
  
          :param batches: a iterator over NHWC numpy arrays in [0, 255].
          :return: a tuple of numpy arrays of shape [N x X], where X is a feature
                   dimension. The tuple is (pool_3, spatial).
          """
          preds = []
          spatial_preds = []
          it = batches if silent else tqdm(batches)
          for batch in it:
              batch = batch.astype(np.float32)
              pred, spatial_pred = self.sess.run(
                  [self.pool_features, self.spatial_features], {self.image_input: batch}
              )
              preds.append(pred.reshape([pred.shape[0], -1]))
              spatial_preds.append(spatial_pred.reshape([spatial_pred.shape[0], -1]))
          return (
              np.concatenate(preds, axis=0),
              np.concatenate(spatial_preds, axis=0),
          )
  
      def read_statistics(
          self, npz_path: str, activations: Tuple[np.ndarray, np.ndarray]
      ) -> Tuple[FIDStatistics, FIDStatistics]:
          obj = np.load(npz_path)
          if "mu" in list(obj.keys()):
              return FIDStatistics(obj["mu"], obj["sigma"]), FIDStatistics(
                  obj["mu_s"], obj["sigma_s"]
              )
          return tuple(self.compute_statistics(x) for x in activations)
  
      def compute_statistics(self, activations: np.ndarray) -> FIDStatistics:
          mu = np.mean(activations, axis=0)
          sigma = np.cov(activations, rowvar=False)
          return FIDStatistics(mu, sigma)
  
      def compute_inception_score(self, activations: np.ndarray, split_size: int = 5000) -> float:
          softmax_out = []
          for i in range(0, len(activations), self.softmax_batch_size):
              acts = activations[i : i + self.softmax_batch_size]
              softmax_out.append(self.sess.run(self.softmax, feed_dict={self.softmax_input: acts}))
          preds = np.concatenate(softmax_out, axis=0)
          # https://github.com/openai/improved-gan/blob/4f5d1ec5c16a7eceb206f42bfc652693601e1d5c/inception_score/model.py#L46
          scores = []
          for i in range(0, len(preds), split_size):
              part = preds[i : i + split_size]
              kl = part * (np.log(part) - np.log(np.expand_dims(np.mean(part, 0), 0)))
              kl = np.mean(np.sum(kl, 1))
              scores.append(np.exp(kl))
          return float(np.mean(scores))
  
      def compute_prec_recall(
          self, activations_ref: np.ndarray, activations_sample: np.ndarray
      ) -> Tuple[float, float]:
          radii_1 = self.manifold_estimator.manifold_radii(activations_ref)
          radii_2 = self.manifold_estimator.manifold_radii(activations_sample)
          pr = self.manifold_estimator.evaluate_pr(
              activations_ref, radii_1, activations_sample, radii_2
          )
          return (float(pr[0][0]), float(pr[1][0]))
  
  
  class ManifoldEstimator:
      """
      A helper for comparing manifolds of feature vectors.
  
      Adapted from https://github.com/kynkaat/improved-precision-and-recall-metric/blob/f60f25e5ad933a79135c783fcda53de30f42c9b9/precision_recall.py#L57
      """
  
      def __init__(
          self,
          session,
          row_batch_size=10000,
          col_batch_size=10000,
          nhood_sizes=(3,),
          clamp_to_percentile=None,
          eps=1e-5,
      ):
          """
          Estimate the manifold of given feature vectors.
  
          :param session: the TensorFlow session.
          :param row_batch_size: row batch size to compute pairwise distances
                                 (parameter to trade-off between memory usage and performance).
          :param col_batch_size: column batch size to compute pairwise distances.
          :param nhood_sizes: number of neighbors used to estimate the manifold.
          :param clamp_to_percentile: prune hyperspheres that have radius larger than
                                      the given percentile.
          :param eps: small number for numerical stability.
          """
          self.distance_block = DistanceBlock(session)
          self.row_batch_size = row_batch_size
          self.col_batch_size = col_batch_size
          self.nhood_sizes = nhood_sizes
          self.num_nhoods = len(nhood_sizes)
          self.clamp_to_percentile = clamp_to_percentile
          self.eps = eps
  
      def warmup(self):
          feats, radii = (
              np.zeros([1, 2048], dtype=np.float32),
              np.zeros([1, 1], dtype=np.float32),
          )
          self.evaluate_pr(feats, radii, feats, radii)
  
      def manifold_radii(self, features: np.ndarray) -> np.ndarray:
          num_images = len(features)
  
          # Estimate manifold of features by calculating distances to k-NN of each sample.
          radii = np.zeros([num_images, self.num_nhoods], dtype=np.float32)
          distance_batch = np.zeros([self.row_batch_size, num_images], dtype=np.float32)
          seq = np.arange(max(self.nhood_sizes) + 1, dtype=np.int32)
  
          for begin1 in range(0, num_images, self.row_batch_size):
              end1 = min(begin1 + self.row_batch_size, num_images)
              row_batch = features[begin1:end1]
  
              for begin2 in range(0, num_images, self.col_batch_size):
                  end2 = min(begin2 + self.col_batch_size, num_images)
                  col_batch = features[begin2:end2]
  
                  # Compute distances between batches.
                  distance_batch[
                      0 : end1 - begin1, begin2:end2
                  ] = self.distance_block.pairwise_distances(row_batch, col_batch)
  
              # Find the k-nearest neighbor from the current batch.
              radii[begin1:end1, :] = np.concatenate(
                  [
                      x[:, self.nhood_sizes]
                      for x in _numpy_partition(distance_batch[0 : end1 - begin1, :], seq, axis=1)
                  ],
                  axis=0,
              )
  
          if self.clamp_to_percentile is not None:
              max_distances = np.percentile(radii, self.clamp_to_percentile, axis=0)
              radii[radii > max_distances] = 0
          return radii
  
      def evaluate(self, features: np.ndarray, radii: np.ndarray, eval_features: np.ndarray):
          """
          Evaluate if new feature vectors are at the manifold.
          """
          num_eval_images = eval_features.shape[0]
          num_ref_images = radii.shape[0]
          distance_batch = np.zeros([self.row_batch_size, num_ref_images], dtype=np.float32)
          batch_predictions = np.zeros([num_eval_images, self.num_nhoods], dtype=np.int32)
          max_realism_score = np.zeros([num_eval_images], dtype=np.float32)
          nearest_indices = np.zeros([num_eval_images], dtype=np.int32)
  
          for begin1 in range(0, num_eval_images, self.row_batch_size):
              end1 = min(begin1 + self.row_batch_size, num_eval_images)
              feature_batch = eval_features[begin1:end1]
  
              for begin2 in range(0, num_ref_images, self.col_batch_size):
                  end2 = min(begin2 + self.col_batch_size, num_ref_images)
                  ref_batch = features[begin2:end2]
  
                  distance_batch[
                      0 : end1 - begin1, begin2:end2
                  ] = self.distance_block.pairwise_distances(feature_batch, ref_batch)
  
              # From the minibatch of new feature vectors, determine if they are in the estimated manifold.
              # If a feature vector is inside a hypersphere of some reference sample, then
              # the new sample lies at the estimated manifold.
              # The radii of the hyperspheres are determined from distances of neighborhood size k.
              samples_in_manifold = distance_batch[0 : end1 - begin1, :, None] <= radii
              batch_predictions[begin1:end1] = np.any(samples_in_manifold, axis=1).astype(np.int32)
  
              max_realism_score[begin1:end1] = np.max(
                  radii[:, 0] / (distance_batch[0 : end1 - begin1, :] + self.eps), axis=1
              )
              nearest_indices[begin1:end1] = np.argmin(distance_batch[0 : end1 - begin1, :], axis=1)
  
          return {
              "fraction": float(np.mean(batch_predictions)),
              "batch_predictions": batch_predictions,
              "max_realisim_score": max_realism_score,
              "nearest_indices": nearest_indices,
          }
  
      def evaluate_pr(
          self,
          features_1: np.ndarray,
          radii_1: np.ndarray,
          features_2: np.ndarray,
          radii_2: np.ndarray,
      ) -> Tuple[np.ndarray, np.ndarray]:
          """
          Evaluate precision and recall efficiently.
  
          :param features_1: [N1 x D] feature vectors for reference batch.
          :param radii_1: [N1 x K1] radii for reference vectors.
          :param features_2: [N2 x D] feature vectors for the other batch.
          :param radii_2: [N x K2] radii for other vectors.
          :return: a tuple of arrays for (precision, recall):
                   - precision: an np.ndarray of length K1
                   - recall: an np.ndarray of length K2
          """
          features_1_status = np.zeros([len(features_1), radii_2.shape[1]], dtype=np.bool)
          features_2_status = np.zeros([len(features_2), radii_1.shape[1]], dtype=np.bool)
          for begin_1 in range(0, len(features_1), self.row_batch_size):
              end_1 = begin_1 + self.row_batch_size
              batch_1 = features_1[begin_1:end_1]
              for begin_2 in range(0, len(features_2), self.col_batch_size):
                  end_2 = begin_2 + self.col_batch_size
                  batch_2 = features_2[begin_2:end_2]
                  batch_1_in, batch_2_in = self.distance_block.less_thans(
                      batch_1, radii_1[begin_1:end_1], batch_2, radii_2[begin_2:end_2]
                  )
                  features_1_status[begin_1:end_1] |= batch_1_in
                  features_2_status[begin_2:end_2] |= batch_2_in
          return (
              np.mean(features_2_status.astype(np.float64), axis=0),
              np.mean(features_1_status.astype(np.float64), axis=0),
          )
  
  
  class DistanceBlock:
      """
      Calculate pairwise distances between vectors.
  
      Adapted from https://github.com/kynkaat/improved-precision-and-recall-metric/blob/f60f25e5ad933a79135c783fcda53de30f42c9b9/precision_recall.py#L34
      """
  
      def __init__(self, session):
          self.session = session
  
          # Initialize TF graph to calculate pairwise distances.
          with session.graph.as_default():
              self._features_batch1 = tf.placeholder(tf.float32, shape=[None, None])
              self._features_batch2 = tf.placeholder(tf.float32, shape=[None, None])
              distance_block_16 = _batch_pairwise_distances(
                  tf.cast(self._features_batch1, tf.float16),
                  tf.cast(self._features_batch2, tf.float16),
              )
              self.distance_block = tf.cond(
                  tf.reduce_all(tf.math.is_finite(distance_block_16)),
                  lambda: tf.cast(distance_block_16, tf.float32),
                  lambda: _batch_pairwise_distances(self._features_batch1, self._features_batch2),
              )
  
              # Extra logic for less thans.
              self._radii1 = tf.placeholder(tf.float32, shape=[None, None])
              self._radii2 = tf.placeholder(tf.float32, shape=[None, None])
              dist32 = tf.cast(self.distance_block, tf.float32)[..., None]
              self._batch_1_in = tf.math.reduce_any(dist32 <= self._radii2, axis=1)
              self._batch_2_in = tf.math.reduce_any(dist32 <= self._radii1[:, None], axis=0)
  
      def pairwise_distances(self, U, V):
          """
          Evaluate pairwise distances between two batches of feature vectors.
          """
          return self.session.run(
              self.distance_block,
              feed_dict={self._features_batch1: U, self._features_batch2: V},
          )
  
      def less_thans(self, batch_1, radii_1, batch_2, radii_2):
          return self.session.run(
              [self._batch_1_in, self._batch_2_in],
              feed_dict={
                  self._features_batch1: batch_1,
                  self._features_batch2: batch_2,
                  self._radii1: radii_1,
                  self._radii2: radii_2,
              },
          )
  
  
  def _batch_pairwise_distances(U, V):
      """
      Compute pairwise distances between two batches of feature vectors.
      """
      with tf.variable_scope("pairwise_dist_block"):
          # Squared norms of each row in U and V.
          norm_u = tf.reduce_sum(tf.square(U), 1)
          norm_v = tf.reduce_sum(tf.square(V), 1)
  
          # norm_u as a column and norm_v as a row vectors.
          norm_u = tf.reshape(norm_u, [-1, 1])
          norm_v = tf.reshape(norm_v, [1, -1])
  
          # Pairwise squared Euclidean distances.
          D = tf.maximum(norm_u - 2 * tf.matmul(U, V, False, True) + norm_v, 0.0)
  
      return D
  
  
  class NpzArrayReader(ABC):
      @abstractmethod
      def read_batch(self, batch_size: int) -> Optional[np.ndarray]:
          pass
  
      @abstractmethod
      def remaining(self) -> int:
          pass
  
      def read_batches(self, batch_size: int) -> Iterable[np.ndarray]:
          def gen_fn():
              while True:
                  batch = self.read_batch(batch_size)
                  if batch is None:
                      break
                  yield batch
  
          rem = self.remaining()
          num_batches = rem // batch_size + int(rem % batch_size != 0)
          return BatchIterator(gen_fn, num_batches)
  
  
  class BatchIterator:
      def __init__(self, gen_fn, length):
          self.gen_fn = gen_fn
          self.length = length
  
      def __len__(self):
          return self.length
  
      def __iter__(self):
          return self.gen_fn()
  
  
  class StreamingNpzArrayReader(NpzArrayReader):
      def __init__(self, arr_f, shape, dtype):
          self.arr_f = arr_f
          self.shape = shape
          self.dtype = dtype
          self.idx = 0
  
      def read_batch(self, batch_size: int) -> Optional[np.ndarray]:
          if self.idx >= self.shape[0]:
              return None
  
          bs = min(batch_size, self.shape[0] - self.idx)
          self.idx += bs
  
          if self.dtype.itemsize == 0:
              return np.ndarray([bs, *self.shape[1:]], dtype=self.dtype)
  
          read_count = bs * np.prod(self.shape[1:])
          read_size = int(read_count * self.dtype.itemsize)
          data = _read_bytes(self.arr_f, read_size, "array data")
          return np.frombuffer(data, dtype=self.dtype).reshape([bs, *self.shape[1:]])
  
      def remaining(self) -> int:
          return max(0, self.shape[0] - self.idx)
  
  
  class MemoryNpzArrayReader(NpzArrayReader):
      def __init__(self, arr):
          self.arr = arr
          self.idx = 0
  
      @classmethod
      def load(cls, path: str, arr_name: str):
          with open(path, "rb") as f:
              arr = np.load(f)[arr_name]
          return cls(arr)
  
      def read_batch(self, batch_size: int) -> Optional[np.ndarray]:
          if self.idx >= self.arr.shape[0]:
              return None
  
          res = self.arr[self.idx : self.idx + batch_size]
          self.idx += batch_size
          return res
  
      def remaining(self) -> int:
          return max(0, self.arr.shape[0] - self.idx)
  
  
  @contextmanager
  def open_npz_array(path: str, arr_name: str) -> NpzArrayReader:
      with _open_npy_file(path, arr_name) as arr_f:
          version = np.lib.format.read_magic(arr_f)
          if version == (1, 0):
              header = np.lib.format.read_array_header_1_0(arr_f)
          elif version == (2, 0):
              header = np.lib.format.read_array_header_2_0(arr_f)
          else:
              yield MemoryNpzArrayReader.load(path, arr_name)
              return
          shape, fortran, dtype = header
          if fortran or dtype.hasobject:
              yield MemoryNpzArrayReader.load(path, arr_name)
          else:
              yield StreamingNpzArrayReader(arr_f, shape, dtype)
  
  
  def _read_bytes(fp, size, error_template="ran out of data"):
      """
      Copied from: https://github.com/numpy/numpy/blob/fb215c76967739268de71aa4bda55dd1b062bc2e/numpy/lib/format.py#L788-L886
  
      Read from file-like object until size bytes are read.
      Raises ValueError if not EOF is encountered before size bytes are read.
      Non-blocking objects only supported if they derive from io objects.
      Required as e.g. ZipExtFile in python 2.6 can return less data than
      requested.
      """
      data = bytes()
      while True:
          # io files (default in python3) return None or raise on
          # would-block, python2 file will truncate, probably nothing can be
          # done about that.  note that regular files can't be non-blocking
          try:
              r = fp.read(size - len(data))
              data += r
              if len(r) == 0 or len(data) == size:
                  break
          except io.BlockingIOError:
              pass
      if len(data) != size:
          msg = "EOF: reading %s, expected %d bytes got %d"
          raise ValueError(msg % (error_template, size, len(data)))
      else:
          return data
  
  
  @contextmanager
  def _open_npy_file(path: str, arr_name: str):
      with open(path, "rb") as f:
          with zipfile.ZipFile(f, "r") as zip_f:
              if f"{arr_name}.npy" not in zip_f.namelist():
                  raise ValueError(f"missing {arr_name} in npz file")
              with zip_f.open(f"{arr_name}.npy", "r") as arr_f:
                  yield arr_f
  
  
  def _download_inception_model():
      if os.path.exists(INCEPTION_V3_PATH):
          return
      print("downloading InceptionV3 model...")
      with requests.get(INCEPTION_V3_URL, stream=True) as r:
          r.raise_for_status()
          tmp_path = INCEPTION_V3_PATH + ".tmp"
          with open(tmp_path, "wb") as f:
              for chunk in tqdm(r.iter_content(chunk_size=8192)):
                  f.write(chunk)
          os.rename(tmp_path, INCEPTION_V3_PATH)
  
  
  def _create_feature_graph(input_batch):
      _download_inception_model()
      prefix = f"{random.randrange(2**32)}_{random.randrange(2**32)}"
      with open(INCEPTION_V3_PATH, "rb") as f:
          graph_def = tf.GraphDef()
          graph_def.ParseFromString(f.read())
      pool3, spatial = tf.import_graph_def(
          graph_def,
          input_map={f"ExpandDims:0": input_batch},
          return_elements=[FID_POOL_NAME, FID_SPATIAL_NAME],
          name=prefix,
      )
      _update_shapes(pool3)
      spatial = spatial[..., :7]
      return pool3, spatial
  
  
  def _create_softmax_graph(input_batch):
      _download_inception_model()
      prefix = f"{random.randrange(2**32)}_{random.randrange(2**32)}"
      with open(INCEPTION_V3_PATH, "rb") as f:
          graph_def = tf.GraphDef()
          graph_def.ParseFromString(f.read())
      (matmul,) = tf.import_graph_def(
          graph_def, return_elements=[f"softmax/logits/MatMul"], name=prefix
      )
      w = matmul.inputs[1]
      logits = tf.matmul(input_batch, w)
      return tf.nn.softmax(logits)
  
  
  def _update_shapes(pool3):
      # https://github.com/bioinf-jku/TTUR/blob/73ab375cdf952a12686d9aa7978567771084da42/fid.py#L50-L63
      ops = pool3.graph.get_operations()
      for op in ops:
          for o in op.outputs:
              shape = o.get_shape()
              if shape._dims is not None:  # pylint: disable=protected-access
                  # shape = [s.value for s in shape] TF 1.x
                  shape = [s for s in shape]  # TF 2.x
                  new_shape = []
                  for j, s in enumerate(shape):
                      if s == 1 and j == 0:
                          new_shape.append(None)
                      else:
                          new_shape.append(s)
                  o.__dict__["_shape_val"] = tf.TensorShape(new_shape)
      return pool3
  
  
  def _numpy_partition(arr, kth, **kwargs):
      num_workers = min(cpu_count(), len(arr))
      chunk_size = len(arr) // num_workers
      extra = len(arr) % num_workers
  
      start_idx = 0
      batches = []
      for i in range(num_workers):
          size = chunk_size + (1 if i < extra else 0)
          batches.append(arr[start_idx : start_idx + size])
          start_idx += size
  
      with ThreadPool(num_workers) as pool:
          return list(pool.map(partial(np.partition, kth=kth, **kwargs), batches))
  
  
  if __name__ == "__main__":
      print(REQUIREMENTS)
      main()
  