code: def _private_compare(
      prot,
      x_bits: PondPrivateTensor,
      r: PondPublicTensor,
      beta: PondPublicTensor,
  ):
      """Logic for private comparison."""
      # TODO[Morten] no need to check this (should be free)
      assert x_bits.backing_dtype == prot.prime_factory
      assert (
          r.backing_dtype.native_type == prot.tensor_factory.native_type
      )
  
      out_shape = r.shape
      out_dtype = r.backing_dtype
      prime_dtype = x_bits.backing_dtype
      bit_length = x_bits.shape[-1]
  
      assert r.shape == out_shape
      assert r.backing_dtype.native_type == out_dtype.native_type
      assert not r.is_scaled
  
      assert x_bits.shape[:-1] == out_shape
      assert x_bits.backing_dtype == prime_dtype
      assert not x_bits.is_scaled
  
      assert beta.shape == out_shape
      assert beta.backing_dtype == prime_dtype
      assert not beta.is_scaled
  
      with tf.name_scope("private_compare"):
          with tf.name_scope("bit_comparisons"):
              # use either r or t = r + 1 according to beta
              s = prot.select(
                  prot.cast_backing(beta, out_dtype), r, r + 1
              )
              s_bits = prot.bits(s, factory=prime_dtype)
              assert s_bits.shape[-1] == bit_length
  
              # compute w_sum
              w_bits = prot.bitwise_xor(x_bits, s_bits)
              w_sum = prot.cumsum(
                  w_bits, axis=-1, reverse=True, exclusive=True
              )
              assert w_sum.backing_dtype == prime_dtype
  
              # compute c, ignoring edge cases at first
              sign = prot.select(beta, 1, -1)
              sign = prot.expand_dims(sign, axis=-1)
              c_except_edge_case = (s_bits - x_bits) * sign + 1 + w_sum
  
              assert c_except_edge_case.backing_dtype == prime_dtype
  
          with tf.name_scope("edge_cases"):
              # adjust for edge cases, i.e. where beta is 1 and s is zero
              # (meaning r was -1)
  
              # identify edge cases
              edge_cases = prot.bitwise_and(
                  beta, prot.equal_zero(s, prime_dtype)
              )
              edge_cases = prot.expand_dims(edge_cases, axis=-1)
  
              # tensor for edge cases: one zero and the rest ones
              c_edge_vals = [0] + [1] * (bit_length - 1)
              c_const = tf.constant(
                  c_edge_vals,
                  dtype=prime_dtype.native_type,
                  shape=(1, bit_length),
              )
              c_edge_case_raw = prime_dtype.tensor(c_const)
              c_edge_case = prot._share_and_wrap(c_edge_case_raw, False)
  
              c = prot.select(
                  edge_cases, c_except_edge_case, c_edge_case
              )
              assert c.backing_dtype == prime_dtype
  
          with tf.name_scope("zero_search"):
              # generate multiplicative mask to hide non-zero values
              with tf.device(prot.server_0.device_name):
                  mask_raw0 = prime_dtype.sample_uniform(
                      c.shape, minval=1
                  )
              with tf.device(prot.server_1.device_name):
                  mask_raw1 = mask_raw0.identity()
              mask = PondPublicTensor(prot, mask_raw0, mask_raw1, False)
  
              # mask non-zero values; this is safe when we're in a prime dtype
              # (since it's a field)
              c_masked = c * mask
              assert c_masked.backing_dtype == prime_dtype
  
              # TODO[Morten] permute
  
              # reconstruct masked values on server 2 to find entries with zeros
              with tf.device(prot.server_2.device_name):
                  d = prot._reconstruct(*c_masked.unwrapped)
                  # find all zero entries
                  zeros = d.equal_zero(out_dtype)
                  # for each bit sequence, determine whether it has one or no zero in it
                  rows_with_zeros = zeros.reduce_sum(
                      axis=-1, keepdims=False
                  )
                  # reshare result
                  result = prot._share_and_wrap(rows_with_zeros, False)
  
          assert (
              result.backing_dtype.native_type == out_dtype.native_type
          )
          return result
  
imports: from __future__ import absolute_import
  import math
  import random
  import sys
  from typing import Optional
  from typing import Tuple
  import numpy
  import tensorflow
  from tf_encrypted.config import get_config
  from tf_encrypted.player import Player
  from tf_encrypted.protocol import memoize
  from tf_encrypted.protocol.pond import Pond
  from tf_encrypted.protocol.pond import PondMaskedTensor
  from tf_encrypted.protocol.pond import PondPrivateTensor
  from tf_encrypted.protocol.pond import PondPublicTensor
  from tf_encrypted.protocol.pond import PondTensor
  from tf_encrypted.protocol.pond import _type
  from tf_encrypted.protocol.protocol import nodes
  from tf_encrypted.protocol.securenn.odd_tensor import oddint64_factory
  from tf_encrypted.tensor import int64factory
  from tf_encrypted.tensor import native_factory
  from tf_encrypted.tensor.factory import AbstractFactory
  from tf_encrypted.tensor.factory import AbstractTensor
label_desc: Use image node_id 5 for calling a global function with example usage: _private_compare(prot, x_bits, r, beta) and returns: result