Python tensorflow.compat.v1.bool() Examples

def _rollout_metadata(batch_env, distributional_size=1):
  """Metadata for rollouts."""
  batch_env_shape = batch_env.observ.get_shape().as_list()
  batch_size = [batch_env_shape[0]]
  value_size = batch_size
  if distributional_size > 1:
    value_size = batch_size + [distributional_size]
  shapes_types_names = [
      # TODO(piotrmilos): possibly retrieve the observation type for batch_env
      (batch_size + batch_env_shape[1:], batch_env.observ_dtype, "observation"),
      (batch_size, tf.float32, "reward"),
      (batch_size, tf.bool, "done"),
      (batch_size + list(batch_env.action_shape), batch_env.action_dtype,
       "action"),
      (batch_size, tf.float32, "pdf"),
      (value_size, tf.float32, "value_function"),
  ]
  return shapes_types_names 

def tracks_own_finished(self):
    """Describes whether the Decoder keeps track of finished states.

    Most decoders will emit a true/false `finished` value independently
    at each time step.  In this case, the `dynamic_decode` function keeps track
    of which batch entries are already finished, and performs a logical OR to
    insert new batches to the finished set.

    Some decoders, however, shuffle batches / beams between time steps and
    `dynamic_decode` will mix up the finished state across these entries because
    it does not track the reshuffle across time steps.  In this case, it is
    up to the decoder to declare that it will keep track of its own finished
    state by setting this property to `True`.

    Returns:
      Python bool.
    """
    return False

  # TODO(scottzhu): Add build/get_config/from_config and other layer methods. 

def get_signature_def(model, use_tf_sampling):
  """Creates a signature def for the SavedModel."""
  if use_tf_sampling:
    return tf.saved_model.signature_def_utils.predict_signature_def(
        inputs={
            'pianorolls': model.inputs['pianorolls'],
        }, outputs={
            'predictions': tf.cast(model.samples, tf.bool),
        })
  return tf.saved_model.signature_def_utils.predict_signature_def(
      inputs={
          'pianorolls': model.model.pianorolls,
          'masks': model.model.masks,
          'lengths': model.model.lengths,
      }, outputs={
          'predictions': model.model.predictions
      }) 

def get_placeholders(self):
    hparams = self.hparams
    return dict(
        pianorolls=tf.placeholder(
            tf.bool,
            [None, None, hparams.num_pitches, hparams.num_instruments],
            "pianorolls"),
        # The default value is only used for checking if completion masker
        # should be evoked.  It can't be used directly as the batch size
        # and length of pianorolls are unknown during static time.
        outer_masks=tf.placeholder_with_default(
            np.zeros(
                (1, 1, hparams.num_pitches, hparams.num_instruments),
                dtype=np.float32),
            [None, None, hparams.num_pitches, hparams.num_instruments],
            "outer_masks"),
        sample_steps=tf.placeholder_with_default(0, (), "sample_steps"),
        total_gibbs_steps=tf.placeholder_with_default(
            0, (), "total_gibbs_steps"),
        current_step=tf.placeholder_with_default(0, (), "current_step"),
        temperature=tf.placeholder_with_default(0.99, (), "temperature")) 

def _flat_reconstruction_loss(self, flat_x_target, flat_rnn_output):
    b_enc, b_dec = tf.split(
        flat_rnn_output,
        [self._nade.num_hidden, self._output_depth], axis=1)
    ll, cond_probs = self._nade.log_prob(
        flat_x_target, b_enc=b_enc, b_dec=b_dec)
    r_loss = -ll
    flat_truth = tf.cast(flat_x_target, tf.bool)
    flat_predictions = tf.greater_equal(cond_probs, 0.5)

    metric_map = {
        'metrics/accuracy':
            tf.metrics.mean(
                tf.reduce_all(tf.equal(flat_truth, flat_predictions), axis=-1)),
        'metrics/recall':
            tf.metrics.recall(flat_truth, flat_predictions),
        'metrics/precision':
            tf.metrics.precision(flat_truth, flat_predictions),
    }

    return r_loss, metric_map 

def _build_replay_buffer(self, use_staging):
    """Creates the replay buffer used by the agent.

    Args:
      use_staging: bool, if True, uses a staging area to prefetch data for
        faster training.

    Returns:
      A WrapperReplayBuffer object.
    """
    return dqn_agent.wrapped_replay_buffer(
        observation_shape=self.observation_shape,
        stack_size=self.stack_size,
        use_staging=use_staging,
        update_horizon=self.update_horizon,
        gamma=self.gamma,
        observation_dtype=self.observation_dtype,
        action_shape=self._env_action_space.shape,
        action_dtype=self._env_action_space.dtype,
        reward_shape=self._response_adapter.response_shape,
        reward_dtype=self._response_adapter.response_dtype) 

def __init__(self, sample_fn, sample_shape, sample_dtype,
               start_inputs, end_fn, next_inputs_fn=None):
    """Initializer.

    Args:
      sample_fn: A callable that takes `outputs` and emits tensor `sample_ids`.
      sample_shape: Either a list of integers, or a 1-D Tensor of type `int32`,
        the shape of the each sample in the batch returned by `sample_fn`.
      sample_dtype: the dtype of the sample returned by `sample_fn`.
      start_inputs: The initial batch of inputs.
      end_fn: A callable that takes `sample_ids` and emits a `bool` vector
        shaped `[batch_size]` indicating whether each sample is an end token.
      next_inputs_fn: (Optional) A callable that takes `sample_ids` and returns
        the next batch of inputs. If not provided, `sample_ids` is used as the
        next batch of inputs.
    """
    self._sample_fn = sample_fn
    self._end_fn = end_fn
    self._sample_shape = tf.TensorShape(sample_shape)
    self._sample_dtype = sample_dtype
    self._next_inputs_fn = next_inputs_fn
    self._batch_size = tf.shape(start_inputs)[0]
    self._start_inputs = tf.convert_to_tensor(
        start_inputs, name="start_inputs") 

def _build(self, inputs, labels):

    def cond(i, unused_attack, success):
      # If we are already successful, we break.
      return tf.logical_and(i < self._num_restarts,
                            tf.logical_not(tf.reduce_all(success)))

    def body(i, attack, success):
      new_attack = self._inner_attack(inputs, labels)
      new_success = self._inner_attack.success
      # The first iteration always sets the attack.
      use_new_values = tf.logical_or(tf.equal(i, 0), new_success)
      return (i + 1,
              tf.where(use_new_values, new_attack, attack),
              tf.logical_or(success, new_success))

    _, self._attack, self._success = tf.while_loop(
        cond, body, back_prop=False, parallel_iterations=1,
        loop_vars=[
            tf.constant(0, dtype=tf.int32),
            inputs,
            tf.zeros([tf.shape(inputs)[0]], dtype=tf.bool),
        ])
    self._logits = self._eval_fn(self._attack, mode='final')
    return self._attack 

def _apply_func_with_prob(func, image, args, prob, bboxes):
  """Apply `func` to image w/ `args` as input with probability `prob`."""
  assert isinstance(args, tuple)
  if six.PY2:
    # pylint: disable=deprecated-method
    arg_spec = inspect.getargspec(func)
    # pylint: enable=deprecated-method
  else:
    arg_spec = inspect.getfullargspec(func)
  assert 'bboxes' == arg_spec[0][1]

  # If prob is a function argument, then this randomness is being handled
  # inside the function, so make sure it is always called.
  if 'prob' in arg_spec[0]:
    prob = 1.0

  # Apply the function with probability `prob`.
  should_apply_op = tf.cast(
      tf.floor(tf.random_uniform([], dtype=tf.float32) + prob), tf.bool)
  augmented_image, augmented_bboxes = tf.cond(
      should_apply_op,
      lambda: func(image, bboxes, *args),
      lambda: (image, bboxes))
  return augmented_image, augmented_bboxes 

def __init__(self, regularizers_to_group):
    """Creates an instance.

    Args:
      regularizers_to_group: A list of generic_regularizers.OpRegularizer
        objects.Their regularization_vector (alive_vector) are expected to be of
        the same length.

    Raises:
      ValueError: regularizers_to_group is not of length at least 2.
    """
    if len(regularizers_to_group) < 2:
      raise ValueError('Groups must be of at least size 2.')
    self._regularization_vector = tf.add_n(
        [r.regularization_vector for r in regularizers_to_group])
    self._alive_vector = tf.cast(
        tf.ones(self._regularization_vector.get_shape()[-1]), tf.bool) 

def test_blackout_pixel_weights_by_box_regions(self):
    def graph_fn():
      boxes = tf.constant(
          [[0.0, 0.0, 5, 5], [0.0, 0.0, 10.0, 20.0], [6.0, 12.0, 8.0, 18.0]],
          dtype=tf.float32)
      blackout = tf.constant([True, False, True], dtype=tf.bool)
      blackout_pixel_weights_by_box_regions = tf.function(
          ta_utils.blackout_pixel_weights_by_box_regions)
      output = blackout_pixel_weights_by_box_regions(10, 20, boxes, blackout)
      return output

    output = self.execute(graph_fn, [])
    # All zeros in region [0:6, 0:6].
    self.assertAlmostEqual(np.sum(output[0:6, 0:6]), 0.0)
    # All zeros in region [12:19, 6:9].
    self.assertAlmostEqual(np.sum(output[6:9, 12:19]), 0.0)
    # All other pixel weights should be 1.0.
    # 20 * 10 - 6 * 6 - 3 * 7 = 143.0
    self.assertAlmostEqual(np.sum(output), 143.0) 

def lengths_to_area_mask(feature_length, length, max_area_size):
  """Generates a non-padding mask for areas based on lengths.

  Args:
    feature_length: a tensor of [batch_size]
    length: the length of the batch
    max_area_size: the maximum area size considered
  Returns:
    mask: a tensor in shape of [batch_size, num_areas]
  """

  paddings = tf.cast(tf.expand_dims(
      tf.logical_not(
          tf.sequence_mask(feature_length, maxlen=length)), 2), tf.float32)
  _, _, area_sum, _, _ = compute_area_features(paddings,
                                               max_area_width=max_area_size)
  mask = tf.squeeze(tf.logical_not(tf.cast(area_sum, tf.bool)), [2])
  return mask 

def test_keypoints(self):
    keypoints = test_utils.keypoints_with_dynamic_shape(10, 16, 4)
    visibilities = tf.cast(tf.random.uniform(tf.shape(keypoints)[:-1], minval=0,
                                             maxval=2, dtype=tf.int32), tf.bool)
    input_tensor_dict = {
        fields.InputDataFields.groundtruth_keypoints:
            test_utils.keypoints_with_dynamic_shape(10, 16, 4),
        fields.InputDataFields.groundtruth_keypoint_visibilities:
            visibilities
    }
    padded_tensor_dict = inputs.pad_input_data_to_static_shapes(
        tensor_dict=input_tensor_dict,
        max_num_boxes=3,
        num_classes=3,
        spatial_image_shape=[5, 6])

    self.assertAllEqual(
        padded_tensor_dict[fields.InputDataFields.groundtruth_keypoints]
        .shape.as_list(), [3, 16, 4])
    self.assertAllEqual(
        padded_tensor_dict[
            fields.InputDataFields.groundtruth_keypoint_visibilities]
        .shape.as_list(), [3, 16]) 

def load_boolq_file(filename, num_par=2):
  """Build a tf.data.Data from a file of boolq examples."""
  tokenizer = tokenization.NltkTokenizer()
  examples = []
  with tf.gfile.Open(filename) as f:
    for line in f:
      obj = json.loads(line)
      context = tokenizer.tokenize(obj["passage"])
      if FLAGS.max_passage_len:
        context = context[:FLAGS.max_passage_len]
      question = tokenizer.tokenize(obj["question"])
      examples.append((question, context, obj["answer"]))

  def get_data():
    out = list(examples)
    np.random.shuffle(out)
    return out

  ds = tf.data.Dataset.from_generator(get_data, (tf.string, tf.string, tf.bool),
                                      ([None], [None], []))

  def to_dict(p, h, label):
    return {"hypothesis": p, "premise": h, "label": label}

  return ds.map(to_dict, num_parallel_calls=num_par) 

def revnet(name, x, hparams, reverse=True):
  """'hparams.depth' steps of generative flow.

  Args:
    name: variable scope for the revnet block.
    x: 4-D Tensor, shape=(NHWC).
    hparams: HParams.
    reverse: bool, forward or backward pass.
  Returns:
    x: 4-D Tensor, shape=(NHWC).
    objective: float.
  """
  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
    steps = np.arange(hparams.depth)
    if reverse:
      steps = steps[::-1]

    objective = 0.0
    for step in steps:
      x, curr_obj = revnet_step(
          "revnet_step_%d" % step, x, hparams, reverse=reverse)
      objective += curr_obj
    return x, objective 

def _minibatch_subsample_fn(self, inputs):
    """Randomly samples anchors for one image.

    Args:
      inputs: a list of 2 inputs. First one is a tensor of shape [num_anchors,
        num_classes] indicating targets assigned to each anchor. Second one
        is a tensor of shape [num_anchors] indicating the class weight of each
        anchor.

    Returns:
      batch_sampled_indicator: bool tensor of shape [num_anchors] indicating
        whether the anchor should be selected for loss computation.
    """
    cls_targets, cls_weights = inputs
    if self._add_background_class:
      # Set background_class bits to 0 so that the positives_indicator
      # computation would not consider background class.
      background_class = tf.zeros_like(tf.slice(cls_targets, [0, 0], [-1, 1]))
      regular_class = tf.slice(cls_targets, [0, 1], [-1, -1])
      cls_targets = tf.concat([background_class, regular_class], 1)
    positives_indicator = tf.reduce_sum(cls_targets, axis=1)
    return self._random_example_sampler.subsample(
        tf.cast(cls_weights, tf.bool),
        batch_size=None,
        labels=tf.cast(positives_indicator, tf.bool)) 

def _padded_batched_proposals_indicator(self,
                                          num_proposals,
                                          max_num_proposals):
    """Creates indicator matrix of non-pad elements of padded batch proposals.

    Args:
      num_proposals: Tensor of type tf.int32 with shape [batch_size].
      max_num_proposals: Maximum number of proposals per image (integer).

    Returns:
      A Tensor of type tf.bool with shape [batch_size, max_num_proposals].
    """
    batch_size = tf.size(num_proposals)
    tiled_num_proposals = tf.tile(
        tf.expand_dims(num_proposals, 1), [1, max_num_proposals])
    tiled_proposal_index = tf.tile(
        tf.expand_dims(tf.range(max_num_proposals), 0), [batch_size, 1])
    return tf.greater(tiled_num_proposals, tiled_proposal_index) 

def batch_boolean_mask(mask):
  """Get indices of true values.

  Args:
    mask: [batch_size, num_values]

  Returns:
    true_indices: [batch_size, max_true]
    gathered_mask: [batch_size, max_true]
  """
  # [batch_size, num_values]
  mask = tf.to_int32(mask)

  # [batch_size]
  num_true = tf.reduce_sum(mask, 1)

  # []
  max_true = tf.reduce_max(num_true)

  # [batch_size, max_true]
  gathered_mask, true_indices = tf.nn.top_k(mask, max_true)
  gathered_mask = tf.cast(gathered_mask, tf.bool)

  return gathered_mask, true_indices 

def pad_batch(features, batch_multiple):
  """Pad batch dim of features to nearest multiple of batch_multiple."""
  feature = list(features.items())[0][1]
  batch_size = tf.shape(feature)[0]
  mod = batch_size % batch_multiple
  has_mod = tf.cast(tf.cast(mod, tf.bool), tf.int32)
  batch_padding = batch_multiple * has_mod - mod

  padded_features = {}
  for k, feature in features.items():
    rank = len(feature.shape)
    paddings = [[0, 0] for _ in range(rank)]
    paddings[0][1] = batch_padding
    padded_feature = tf.pad(feature, paddings)
    padded_features[k] = padded_feature
  return padded_features


# TODO(lukaszkaiser): refactor the API to not be just a list of self params
#   but make sense for other uses too. 

def set_recall(predictions, labels, weights_fn=common_layers.weights_nonzero):
  """Recall of set predictions.

  Args:
    predictions : A Tensor of scores of shape [batch, nlabels].
    labels: A Tensor of int32s giving true set elements,
      of shape [batch, seq_length].
    weights_fn: A function to weight the elements.

  Returns:
    hits: A Tensor of shape [batch, nlabels].
    weights: A Tensor of shape [batch, nlabels].
  """
  with tf.variable_scope("set_recall", values=[predictions, labels]):
    labels = tf.squeeze(labels, [2, 3])
    weights = weights_fn(labels)
    labels = tf.one_hot(labels, predictions.shape[-1])
    labels = tf.reduce_max(labels, axis=1)
    labels = tf.cast(labels, tf.bool)
    return tf.to_float(tf.equal(labels, predictions)), weights 

def sparse_reduce(sp_tensor, rank, agg_fn="sum", axis=-1):
  """Reduce SparseTensor along the given axis.

  Args:
    sp_tensor: SparseTensor of arbitrary rank.
    rank: Integer rank of the sparse tensor.
    agg_fn: Reduce function for aggregation.
    axis: Integer specifying axis to sum over.

  Returns:
    sp_tensor: SparseTensor of one less rank.
  """
  if axis < 0:
    axis = rank + axis
  axes_to_keep = tf.one_hot(
      axis, rank, on_value=False, off_value=True, dtype=tf.bool)
  indices_to_keep = tf.boolean_mask(sp_tensor.indices, axes_to_keep, axis=1)
  new_shape = tf.boolean_mask(sp_tensor.dense_shape, axes_to_keep)
  indices_to_keep.set_shape([None, rank - 1])
  indices, values = aggregate_sparse_indices(
      indices_to_keep, sp_tensor.values, new_shape, agg_fn=agg_fn)
  return tf.sparse.reorder(tf.SparseTensor(indices, values, new_shape)) 

def match(self, similarity_matrix, valid_rows=None, scope=None):
    """Computes matches among row and column indices and returns the result.

    Computes matches among the row and column indices based on the similarity
    matrix and optional arguments.

    Args:
      similarity_matrix: Float tensor of shape [N, M] with pairwise similarity
        where higher value means more similar.
      valid_rows: A boolean tensor of shape [N] indicating the rows that are
        valid for matching.
      scope: Op scope name. Defaults to 'Match' if None.

    Returns:
      A Match object with the results of matching.
    """
    with tf.name_scope(scope, 'Match') as scope:
      if valid_rows is None:
        valid_rows = tf.ones(tf.shape(similarity_matrix)[0], dtype=tf.bool)
      return Match(self._match(similarity_matrix, valid_rows),
                   self._use_matmul_gather) 

def testReturnsCorrectWeightedLossWithLossesMask(self):
    batch_size = 4
    num_anchors = 10
    code_size = 4
    def graph_fn():
      prediction_tensor = tf.ones([batch_size, num_anchors, code_size])
      target_tensor = tf.zeros([batch_size, num_anchors, code_size])
      weights = tf.constant([[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
                             [1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
                             [1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
                             [1, 1, 1, 1, 1, 0, 0, 0, 0, 0]], tf.float32)
      losses_mask = tf.constant([True, False, True, True], tf.bool)
      loss_op = losses.WeightedL2LocalizationLoss()
      loss = tf.reduce_sum(loss_op(prediction_tensor, target_tensor,
                                   weights=weights, losses_mask=losses_mask))
      return loss
    expected_loss = (3 * 5 * 4) / 2.0
    loss_output = self.execute(graph_fn, [])
    self.assertAllClose(loss_output, expected_loss) 

def test_set_keypoint_visibilities(self):
    keypoints_np = np.array(
        [
            [[np.nan, 0.2],
             [np.nan, np.nan],
             [-3., 7.]],
            [[0.5, 0.2],
             [4., 1.0],
             [-3., np.nan]],
        ], dtype=np.float32)
    initial_keypoint_visibilities_np = np.array(
        [
            [False,
             True,  # Will be overriden by NaN coords.
             False],  # Will be maintained, even though non-NaN coords.
            [True,
             False,  # Will be maintained, even though non-NaN coords.
             False]
        ])
    def graph_fn():
      keypoints = tf.constant(keypoints_np, dtype=tf.float32)
      initial_keypoint_visibilities = tf.constant(
          initial_keypoint_visibilities_np, dtype=tf.bool)
      keypoint_visibilities = keypoint_ops.set_keypoint_visibilities(
          keypoints, initial_keypoint_visibilities)
      return keypoint_visibilities

    expected_kpt_vis = [
        [False, False, False],
        [True, False, False]
    ]
    output = self.execute(graph_fn, [])
    self.assertAllEqual(expected_kpt_vis, output) 

def reduce_concat_tokens(dataset,
                         feature_key='targets',
                         batch_size=128,
                         **unused_kwargs):
  """Token-preprocessor to concatenate multiple unrelated documents.

  If we want to generate examples of exactly the right length,
  (to avoid wasting space on padding), then we use this function, folowed by
  split_tokens.

  Args:
    dataset: a tf.data.Dataset with dictionaries containing the key feature_key.
    feature_key: an string
    batch_size: an integer - how many documents to concatenate into one

  Returns:
    a dataset
  """
  dataset = dataset.map(lambda x: {feature_key: x[feature_key]},
                        num_parallel_calls=tf.data.experimental.AUTOTUNE)
  dataset = dataset.padded_batch(batch_size, padded_shapes={feature_key: [-1]})
  def _my_fn(x):
    tokens = tf.reshape(x[feature_key], [-1])
    # strip padding
    tokens = tf.boolean_mask(tokens, tf.cast(tokens, tf.bool))
    return {feature_key: tokens}
  return dataset.map(_my_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE) 

def add_sync_queues_and_barrier(self, name_prefix, enqueue_after_list):
    """Adds ops to enqueue on all worker queues.

    Args:
      name_prefix: prefixed for the shared_name of ops.
      enqueue_after_list: control dependency from ops.

    Returns:
      An op that should be used as control dependency before starting next step.
    """
    self.sync_queue_counter += 1
    with tf.device(self.sync_queue_devices[(
        self.sync_queue_counter % len(self.sync_queue_devices))]):
      sync_queues = [
          tf.FIFOQueue(self.num_workers, [tf.bool], shapes=[[]],
                       shared_name='%s%s' % (name_prefix, i))
          for i in range(self.num_workers)]
      queue_ops = []
      # For each other worker, add an entry in a queue, signaling that it can
      # finish this step.
      token = tf.constant(False)
      with tf.control_dependencies(enqueue_after_list):
        for i, q in enumerate(sync_queues):
          if i == self.task_index:
            queue_ops.append(tf.no_op())
          else:
            queue_ops.append(q.enqueue(token))

      # Drain tokens off queue for this worker, one for each other worker.
      queue_ops.append(
          sync_queues[self.task_index].dequeue_many(len(sync_queues) - 1))

      return tf.group(*queue_ops) 

def create_cpc_model_and_placeholders(num_choices):
  """Build model and placeholders."""

  bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config)
  is_training = False
  use_one_hot_embeddings = False
  seq_length = 512

  Placeholders = namedtuple("Placeholders", [
      "input_ids", "input_mask", "segment_ids", "labels", "label_types",
      "softmax_mask"
  ])

  input_ids = tf.placeholder(dtype=tf.int32, shape=[None, seq_length])
  input_mask = tf.placeholder(dtype=tf.int32, shape=[None, seq_length])
  segment_ids = tf.placeholder(dtype=tf.int32, shape=[None, seq_length])
  labels = tf.placeholder(dtype=tf.int32, shape=[None, 8])
  label_types = tf.placeholder(dtype=tf.int32, shape=[None, 8])
  softmax_mask = tf.placeholder(dtype=tf.bool, shape=[None])
  placeholders = Placeholders(input_ids, input_mask, segment_ids, labels,
                              label_types, softmax_mask)

  model = modeling.BertModel(
      config=bert_config,
      is_training=is_training,
      input_ids=input_ids,
      input_mask=input_mask,
      token_type_ids=segment_ids,
      use_one_hot_embeddings=use_one_hot_embeddings)

  logits, probabilities = create_cpc_model(model, num_choices, False,
                                           softmax_mask)

  Model = namedtuple("Model", ["logits", "probabilities"])
  model = Model(logits, probabilities)

  return placeholders, model 

def select_slate_optimal(slate_size, s_no_click, s, q):
  """Selects the slate using exhaustive search.

  This algorithm corresponds to the method "OS" in
  Ie et al. https://arxiv.org/abs/1905.12767.

  Args:
    slate_size: int, the size of the recommendation slate.
    s_no_click: float tensor, the score for not clicking any document.
    s: [num_of_documents] tensor, the scores for clicking documents.
    q: [num_of_documents] tensor, the predicted q values for documents.

  Returns:
    [slate_size] tensor, the selected slate.
  """

  num_candidates = s.shape.as_list()[0]

  # Obtain all possible slates given current docs in the candidate set.
  mesh_args = [list(range(num_candidates))] * slate_size
  slates = tf.stack(tf.meshgrid(*mesh_args), axis=-1)
  slates = tf.reshape(slates, shape=(-1, slate_size))

  # Filter slates that include duplicates to ensure each document is picked
  # at most once.
  unique_mask = tf.map_fn(
      lambda x: tf.equal(tf.size(input=x), tf.size(input=tf.unique(x)[0])),
      slates,
      dtype=tf.bool)
  slates = tf.boolean_mask(tensor=slates, mask=unique_mask)

  slate_q_values = tf.gather(s * q, slates)
  slate_scores = tf.gather(s, slates)
  slate_normalizer = tf.reduce_sum(
      input_tensor=slate_scores, axis=1) + s_no_click

  slate_q_values = slate_q_values / tf.expand_dims(slate_normalizer, 1)
  slate_sum_q_values = tf.reduce_sum(input_tensor=slate_q_values, axis=1)
  max_q_slate_index = tf.argmax(input=slate_sum_q_values)
  return tf.gather(slates, max_q_slate_index, axis=0) 

def sequence_mask(lengths, maxlen=None, dtype=tf.bool, transpose=False):
  """Returns a sequence mask, like tf.sequence_mask().

  Unlike tf.sequence_mask(), this can also generate a transposed mask, which is
  convenient for working with time-major sequences.

  Args:
    lengths: <int>[batch_size] Sequence lengths.
    maxlen: <int>[] Maximum length, or None to compute the max of `lengths`.
    dtype: DType of the generated mask.
    transpose: Whether to generate the transpose of the mask.

  Returns:
    <dtype>[maxlen, batch_size] Sequence mask.
  """
  with tf.name_scope("sequence_mask"):
    if maxlen is None:
      maxlen = tf.reduce_max(lengths)
    positions = tf.range(maxlen, dtype=lengths.dtype)

    positions_singleton_axis = 1 if transpose else 0
    lengths_singleton_axis = 0 if transpose else 1
    positions = tf.expand_dims(positions, axis=positions_singleton_axis)
    lengths = tf.expand_dims(lengths, axis=lengths_singleton_axis)
    mask = positions < lengths

    if dtype != tf.bool:
      mask = tf.cast(mask, dtype)
    return mask 

def trues_like(t):
  return tf.cast(tf.ones_like(t), tf.bool)