jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _cat_sample(self, sample_shape=(), seed=None, name="sample"):
    assert len(sample_shape) in (0, 1), sample_shape
    assert len(self.logits_parameter().shape) == 2
    indices = tf.random.categorical(
        self.logits_parameter(), sample_shape[0] if sample_shape else 1, self.dtype, seed, name
    )
    if not sample_shape:
        indices = indices[..., 0]
    return indices
//  [python]
def create_variables(self):

    with tf.name_scope("model_inputs"):
      # raw state representation
      self.states = tf.placeholder(tf.float32, (None, self.state_dim), name="states")

    # rollout action based on current policy
    with tf.name_scope("predict_actions"):
      # initialize policy network
      with tf.variable_scope("policy_network"):
        self.policy_outputs = self.policy_network(self.states)

      # predict actions from policy network
      self.action_scores = tf.identity(self.policy_outputs, name="action_scores")
      # Note 1: tf.multinomial is not good enough to use yet
      # so we don't use self.predicted_actions for now
      self.predicted_actions = tf.multinomial(self.action_scores, 1)

    # regularization loss
    policy_network_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope="policy_network")

    # compute loss and gradients
    with tf.name_scope("compute_pg_gradients"):
      # gradients for selecting action from policy network
      self.taken_actions = tf.placeholder(tf.int32, (None,), name="taken_actions")
      self.discounted_rewards = tf.placeholder(tf.float32, (None,), name="discounted_rewards")

      with tf.variable_scope("policy_network", reuse=True):
        self.logprobs = self.policy_network(self.states)

      # compute policy loss and regularization loss
      self.cross_entropy_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=self.logprobs, labels=self.taken_actions)
      self.pg_loss            = tf.reduce_mean(self.cross_entropy_loss)
      self.reg_loss           = tf.reduce_sum([tf.reduce_sum(tf.square(x)) for x in policy_network_variables])
      self.loss               = self.pg_loss + self.reg_param * self.reg_loss

      # compute gradients
      self.gradients = self.optimizer.compute_gradients(self.loss)

      # compute policy gradients
      for i, (grad, var) in enumerate(self.gradients):
        if grad is not None:
          self.gradients[i] = (grad * self.discounted_rewards, var)

      for grad, var in self.gradients:
        tf.summary.histogram(var.name, var)
        if grad is not None:
          tf.summary.histogram(var.name + '/gradients', grad)

      # emit summaries
      tf.summary.scalar("policy_loss", self.pg_loss)
      tf.summary.scalar("reg_loss", self.reg_loss)
      tf.summary.scalar("total_loss", self.loss)

    # training update
    with tf.name_scope("train_policy_network"):
      # apply gradients to update policy network
      self.train_op = self.optimizer.apply_gradients(self.gradients)

    self.summarize = tf.summary.merge_all()
    self.no_op = tf.no_op()
//  [python]
def create_variables(self):

    with tf.name_scope("model_inputs"):
      # raw state representation
      self.states = tf.placeholder(tf.float32, (None, self.state_dim), name="states")

    # rollout action based on current policy
    with tf.name_scope("predict_actions"):
      # initialize policy network
      with tf.variable_scope("policy_network"):
        self.policy_outputs = self.policy_network(self.states)

      # predict actions from policy network
      self.action_scores = tf.identity(self.policy_outputs, name="action_scores")
      # Note 1: tf.multinomial is not good enough to use yet
      # so we don't use self.predicted_actions for now
      self.predicted_actions = tf.random.categorical(self.action_scores, 1)

    # regularization loss
    policy_network_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope="policy_network")

    # compute loss and gradients
    with tf.name_scope("compute_pg_gradients"):
      # gradients for selecting action from policy network
      self.taken_actions = tf.placeholder(tf.int32, (None,), name="taken_actions")
      self.discounted_rewards = tf.placeholder(tf.float32, (None,), name="discounted_rewards")

      with tf.variable_scope("policy_network", reuse=True):
        self.logprobs = self.policy_network(self.states)

      # compute policy loss and regularization loss
      self.cross_entropy_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=self.logprobs, labels=self.taken_actions)
      self.pg_loss            = tf.reduce_mean(self.cross_entropy_loss)
      self.reg_loss           = tf.reduce_sum([tf.reduce_sum(tf.square(x)) for x in policy_network_variables])
      self.loss               = self.pg_loss + self.reg_param * self.reg_loss

      # compute gradients
      self.gradients = self.optimizer.compute_gradients(self.loss)

      # compute policy gradients
      for i, (grad, var) in enumerate(self.gradients):
        if grad is not None:
          self.gradients[i] = (grad * self.discounted_rewards, var)

      for grad, var in self.gradients:
        tf.summary.histogram(var.name, var)
        if grad is not None:
          tf.summary.histogram(var.name + '/gradients', grad)

      # emit summaries
      tf.summary.scalar("policy_loss", self.pg_loss)
      tf.summary.scalar("reg_loss", self.reg_loss)
      tf.summary.scalar("total_loss", self.loss)

    # training update
    with tf.name_scope("train_policy_network"):
      # apply gradients to update policy network
      self.train_op = self.optimizer.apply_gradients(self.gradients)

    self.summarize = tf.summary.merge_all()
    self.no_op = tf.no_op()
//  [python]
def __init__(self, name: str, env):
        """
        :param name: string
        :param env: gym env
        """
        ob_space = env.observation_space
        act_space = env.action_space

        with tf.variable_scope(name):
            self.obs = tf.placeholder(dtype=tf.float32, shape=[None] + list(ob_space.shape), name='obs')
            # Actor (Policy): Given a state (or observation)
            # obtain the distribution of actions
            with tf.variable_scope('policy_net'):
                layer_1 = tf.layers.dense(inputs=self.obs, units=20, activation=tf.tanh)
                layer_2 = tf.layers.dense(inputs=layer_1, units=20, activation=tf.tanh)
                layer_3 = tf.layers.dense(inputs=layer_2, units=act_space.n, activation=tf.tanh)
                self.act_probs = tf.layers.dense(inputs=layer_3, units=act_space.n, activation=tf.nn.softmax)

            # Critic 
            with tf.variable_scope('value_net'):
                layer_1 = tf.layers.dense(inputs=self.obs, units=20, activation=tf.tanh)
                layer_2 = tf.layers.dense(inputs=layer_1, units=20, activation=tf.tanh)
                self.v_preds = tf.layers.dense(inputs=layer_2, units=1, activation=None)

            self.act_stochastic = tf.multinomial(tf.log(self.act_probs), num_samples=1)
            self.act_stochastic = tf.reshape(self.act_stochastic, shape=[-1])

            self.act_deterministic = tf.argmax(self.act_probs, axis=1)

            # 辅助功能
            self.scope = tf.get_variable_scope().name
//  [python]
def __init__(self, name: str, env):
        """
        :param name: string
        :param env: gym env
        """
        ob_space = env.observation_space
        act_space = env.action_space

        with tf.variable_scope(name):
            self.obs = tf.placeholder(dtype=tf.float32, shape=[None] + list(ob_space.shape), name='obs')
            # Actor (Policy): Given a state (or observation)
            # obtain the distribution of actions
            with tf.variable_scope('policy_net'):
                layer_1 = tf.layers.dense(inputs=self.obs, units=20, activation=tf.tanh)
                layer_2 = tf.layers.dense(inputs=layer_1, units=20, activation=tf.tanh)
                layer_3 = tf.layers.dense(inputs=layer_2, units=act_space.n, activation=tf.tanh)
                self.act_probs = tf.layers.dense(inputs=layer_3, units=act_space.n, activation=tf.nn.softmax)

            # Critic 
            with tf.variable_scope('value_net'):
                layer_1 = tf.layers.dense(inputs=self.obs, units=20, activation=tf.tanh)
                layer_2 = tf.layers.dense(inputs=layer_1, units=20, activation=tf.tanh)
                self.v_preds = tf.layers.dense(inputs=layer_2, units=1, activation=None)

            self.act_stochastic = tf.random.categorical(tf.log(self.act_probs), num_samples=1)
            self.act_stochastic = tf.reshape(self.act_stochastic, shape=[-1])

            self.act_deterministic = tf.argmax(self.act_probs, axis=1)

            # 辅助功能
            self.scope = tf.get_variable_scope().name
//  [python]
def random_rotate90_input_and_output_images(
    input_img,
    output_img,
    azimuth_in_channel = None,
    azimuth_out_channel = None,
):
  """Randomly rotate input and output images in increments of 90.

  Args:
    input_img: Tensor with dimensions HWC or BHWC.
    output_img: Tensor with dimensions HWC or BHWC.
    azimuth_in_channel: Channel index of azimuth channel in input. Default value
      `None` means there is no azimuth channel.
    azimuth_out_channel: Channel index of azimuth channel in output. Default
      value `None` means there is no azimuth channel.

  Returns:
    input_img: Tensor with dimensions HWC or BHWC.
    output_img: Tensor with dimensions HWC or BHWC.
  """
  num_rotations = tf.multinomial(tf.math.log([[0.25, 0.25, 0.25, 0.25]]), 1)[0][0]
  input_img = tf.image.rot90(input_img, k=num_rotations)
  output_img = tf.image.rot90(output_img, k=num_rotations)

  if azimuth_in_channel:
    input_img = _update_azimuth_rotate90(input_img, azimuth_in_channel,
                                         num_rotations)
  if azimuth_out_channel:
    output_img = _update_azimuth_rotate90(output_img, azimuth_out_channel,
                                          num_rotations)

  return input_img, output_img
//  [python]
def random_rotate90_input_and_output_images(
    input_img,
    output_img,
    azimuth_in_channel = None,
    azimuth_out_channel = None,
):
  """Randomly rotate input and output images in increments of 90.

  Args:
    input_img: Tensor with dimensions HWC or BHWC.
    output_img: Tensor with dimensions HWC or BHWC.
    azimuth_in_channel: Channel index of azimuth channel in input. Default value
      `None` means there is no azimuth channel.
    azimuth_out_channel: Channel index of azimuth channel in output. Default
      value `None` means there is no azimuth channel.

  Returns:
    input_img: Tensor with dimensions HWC or BHWC.
    output_img: Tensor with dimensions HWC or BHWC.
  """
  num_rotations = tf.random.categorical(
      tf.math.log([[0.25, 0.25, 0.25, 0.25]]), 1, dtype='int32')[0][0]
  input_img = tf.image.rot90(input_img, k=num_rotations)
  output_img = tf.image.rot90(output_img, k=num_rotations)

  if azimuth_in_channel:
    input_img = _update_azimuth_rotate90(input_img, azimuth_in_channel,
                                         num_rotations)
  if azimuth_out_channel:
    output_img = _update_azimuth_rotate90(output_img, azimuth_out_channel,
                                          num_rotations)

  return input_img, output_img
//  [python]
def _top_k_sample(logits, ignore_ids=None, num_samples=1, k=10):
    """
    Does top-k sampling. if ignore_ids is on, then we will zero out those logits.
    :param logits: [batch_size, vocab_size] tensor
    :param ignore_ids: [vocab_size] one-hot representation of the indices we'd like to ignore and never predict,
                        like padding maybe
    :param p: topp threshold to use, either a float or a [batch_size] vector
    :return: [batch_size, num_samples] samples
    # TODO FIGURE OUT HOW TO DO THIS ON TPUS. IT'S HELLA SLOW RIGHT NOW, DUE TO ARGSORT I THINK
    """
    with tf.variable_scope('top_p_sample'):
        batch_size, vocab_size = get_shape_list(logits, expected_rank=2)

        probs = tf.nn.softmax(logits if ignore_ids is None else logits - tf.cast(ignore_ids[None], tf.float32) * 1e10,
                              axis=-1)
        # [batch_size, vocab_perm]
        indices = tf.argsort(probs, direction='DESCENDING')

        # find the top pth index to cut off. careful we don't want to cutoff everything!
        # result will be [batch_size, vocab_perm]
        k_expanded = k if isinstance(k, int) else k[:, None]
        exclude_mask = tf.range(vocab_size)[None] >= k_expanded

        # OPTION A - sample in the sorted space, then unsort.
        logits_to_use = tf.batch_gather(logits, indices) - tf.cast(exclude_mask, tf.float32) * 1e10
        sample_perm = tf.multinomial(logits=logits_to_use, num_samples=num_samples)
        sample = tf.batch_gather(indices, sample_perm)

    return {
        'probs': probs,
        'sample': sample,
    }
//  [python]
def _top_k_sample(logits, ignore_ids=None, num_samples=1, k=10):
    """
    Does top-k sampling. if ignore_ids is on, then we will zero out those logits.
    :param logits: [batch_size, vocab_size] tensor
    :param ignore_ids: [vocab_size] one-hot representation of the indices we'd like to ignore and never predict,
                        like padding maybe
    :param p: topp threshold to use, either a float or a [batch_size] vector
    :return: [batch_size, num_samples] samples
    # TODO FIGURE OUT HOW TO DO THIS ON TPUS. IT'S HELLA SLOW RIGHT NOW, DUE TO ARGSORT I THINK
    """
    with tf.variable_scope('top_p_sample'):
        batch_size, vocab_size = get_shape_list(logits, expected_rank=2)

        probs = tf.nn.softmax(logits if ignore_ids is None else logits - tf.cast(ignore_ids[None], tf.float32) * 1e10,
                              axis=-1)
        # [batch_size, vocab_perm]
        indices = tf.argsort(probs, direction='DESCENDING')

        # find the top pth index to cut off. careful we don't want to cutoff everything!
        # result will be [batch_size, vocab_perm]
        k_expanded = k if isinstance(k, int) else k[:, None]
        exclude_mask = tf.range(vocab_size)[None] >= k_expanded

        # OPTION A - sample in the sorted space, then unsort.
        logits_to_use = tf.batch_gather(logits, indices) - tf.cast(exclude_mask, tf.float32) * 1e10
        sample_perm = tf.random.categorical(logits=logits_to_use, num_samples=num_samples)
        sample = tf.batch_gather(indices, sample_perm)

    return {
        'probs': probs,
        'sample': sample,
    }
//  [python]
def _top_k_sample(logits, ignore_ids=None, num_samples=1, k=10):
    """
    Does top-k sampling. if ignore_ids is on, then we will zero out those logits.
    :param logits: [batch_size, vocab_size] tensor
    :param ignore_ids: [vocab_size] one-hot representation of the indices we'd like to ignore and never predict,
                        like padding maybe
    :param p: topp threshold to use, either a float or a [batch_size] vector
    :return: [batch_size, num_samples] samples

    # TODO FIGURE OUT HOW TO DO THIS ON TPUS. IT'S HELLA SLOW RIGHT NOW, DUE TO ARGSORT I THINK
    """
    with tf.variable_scope('top_p_sample'):
        batch_size, vocab_size = get_shape_list(logits, expected_rank=2)

        probs = tf.nn.softmax(logits if ignore_ids is None else logits - tf.cast(ignore_ids[None], tf.float32) * 1e10,
                              axis=-1)
        # [batch_size, vocab_perm]
        indices = tf.argsort(probs, direction='DESCENDING')

        # find the top pth index to cut off. careful we don't want to cutoff everything!
        # result will be [batch_size, vocab_perm]
        k_expanded = k if isinstance(k, int) else k[:, None]
        exclude_mask = tf.range(vocab_size)[None] >= k_expanded

        # OPTION A - sample in the sorted space, then unsort.
        logits_to_use = tf.batch_gather(logits, indices) - tf.cast(exclude_mask, tf.float32) * 1e10
        sample_perm = tf.multinomial(logits=logits_to_use, num_samples=num_samples)
        sample = tf.batch_gather(indices, sample_perm)

    return {
        'probs': probs,
        'sample': sample,
    }
//  [python]
def _top_k_sample(logits, ignore_ids=None, num_samples=1, k=10):
    """
    Does top-k sampling. if ignore_ids is on, then we will zero out those logits.
    :param logits: [batch_size, vocab_size] tensor
    :param ignore_ids: [vocab_size] one-hot representation of the indices we'd like to ignore and never predict,
                        like padding maybe
    :param p: topp threshold to use, either a float or a [batch_size] vector
    :return: [batch_size, num_samples] samples

    # TODO FIGURE OUT HOW TO DO THIS ON TPUS. IT'S HELLA SLOW RIGHT NOW, DUE TO ARGSORT I THINK
    """
    with tf.variable_scope('top_p_sample'):
        batch_size, vocab_size = get_shape_list(logits, expected_rank=2)

        probs = tf.nn.softmax(logits if ignore_ids is None else logits - tf.cast(ignore_ids[None], tf.float32) * 1e10,
                              axis=-1)
        # [batch_size, vocab_perm]
        indices = tf.argsort(probs, direction='DESCENDING')

        # find the top pth index to cut off. careful we don't want to cutoff everything!
        # result will be [batch_size, vocab_perm]
        k_expanded = k if isinstance(k, int) else k[:, None]
        exclude_mask = tf.range(vocab_size)[None] >= k_expanded

        # OPTION A - sample in the sorted space, then unsort.
        logits_to_use = tf.batch_gather(logits, indices) - tf.cast(exclude_mask, tf.float32) * 1e10
        sample_perm = tf.random.categorical(logits=logits_to_use, num_samples=num_samples)
        sample = tf.batch_gather(indices, sample_perm)

    return {
        'probs': probs,
        'sample': sample,
    }
//  [python]
def sample(self, num_samples: int, seed: Optional[int] = None) -> TensorType:
        """
        :param num_samples: The number of points to sample from this search space.
        :param seed: Random seed for reproducibility.
        :return: ``num_samples`` i.i.d. random points, sampled uniformly,
            from this search space.
        """
        if seed is not None:  # ensure reproducibility
            tf.random.set_seed(seed)

        if num_samples == 0:
            return self.points[:0, :]
        else:
            sampled_indices = tf.multinomial(
                tf.ones((1, tf.shape(self.points)[0])), num_samples, seed=seed
            )
            return tf.gather(self.points, sampled_indices)[0, :, :]
//  [python]
def sample(self, num_samples: int, seed: Optional[int] = None) -> TensorType:
        """
        :param num_samples: The number of points to sample from this search space.
        :param seed: Random seed for reproducibility.
        :return: ``num_samples`` i.i.d. random points, sampled uniformly,
            from this search space.
        """
        if seed is not None:  # ensure reproducibility
            tf.random.set_seed(seed)

        if num_samples == 0:
            return self.points[:0, :]
        else:
            sampled_indices = tf.random.categorical(
                tf.ones((1, tf.shape(self.points)[0])), num_samples, seed=seed
            )
            return tf.gather(self.points, sampled_indices)[0, :, :]
//  [python]
def _cat_sample(self, sample_shape=(), seed=None, name='sample'):
    assert len(sample_shape) in (0, 1), sample_shape
    assert len(self.logits_parameter().shape) == 2
    indices = tf.random.categorical(
        self.logits_parameter(), sample_shape[0] if sample_shape else 1,
        self.dtype, seed, name)
    if not sample_shape:
        indices = indices[..., 0]
    return indices
//  [python]
def random_cat():
    return tf.random.categorical(tf.math.log([[0.5, 0.5]]), 5)
//  [python]
def embedding_to_cluster_centers_given_scores(embedding,
                                              centerness_scores,
                                              num_samples_per_center=400):
  """Initializes centers given centerness scores.

  Args:
    embedding: A tf.float32 tensor of size [N, embedding_size].
    centerness_scores: A tf.float32 tensor of size [N]. it is
                       assumed that centerness scores are between 0 and 1.
    num_samples_per_center: Number of input embedding samples per center.

  Returns:
    embedding_seeds: A tf.float32 tensor of size [k, embedding_size].
    indices: Indices of embedding_seeds
  """
  n = tf.shape(embedding)[0]
  num_centers = n // num_samples_per_center
  center_indices = tf.squeeze(
      tf.multinomial(
          tf.math.log(
              tf.expand_dims(centerness_scores, axis=0) + sys.float_info.min),
          num_centers),
      axis=0)
  return tf.gather(embedding, center_indices), center_indices
//  [python]
def embedding_to_cluster_centers_given_scores(embedding,
                                              centerness_scores,
                                              num_samples_per_center=400):
  """Initializes centers given centerness scores.

  Args:
    embedding: A tf.float32 tensor of size [N, embedding_size].
    centerness_scores: A tf.float32 tensor of size [N]. it is
                       assumed that centerness scores are between 0 and 1.
    num_samples_per_center: Number of input embedding samples per center.

  Returns:
    embedding_seeds: A tf.float32 tensor of size [k, embedding_size].
    indices: Indices of embedding_seeds
  """
  n = tf.shape(embedding)[0]
  num_centers = n // num_samples_per_center
  center_indices = tf.squeeze(
      tf.random.categorical(
          tf.math.log(
              tf.expand_dims(centerness_scores, axis=0) + sys.float_info.min),
          num_centers),
      axis=0)
  return tf.gather(embedding, center_indices), center_indices
//  [python]
def sample_actions(logits):
    """ Sample action from unnormalized log probabilities """
    return tf.reshape(tf.random.categorical(logits, num_samples=1), [-1])
//  [python]
def act(self, observation, maze_map):
    observations = tf.convert_to_tensor([observation])
    maze_maps = tf.convert_to_tensor([maze_map])
    states = self.process_states(observations, maze_maps)
    embed = self.embed_states(states)
    logit = self._policy(embed)
    return tf.multinomial(logit, 1, dtype=tf.int32)[:, 0]
//  [python]
def act(self, observation, maze_map):
    observations = tf.convert_to_tensor([observation])
    maze_maps = tf.convert_to_tensor([maze_map])
    states = self.process_states(observations, maze_maps)
    embed = self.embed_states(states)
    logit = self._policy(embed)
    return tf.random.categorical(logit, 1, dtype=tf.int32)[:, 0]
//  [python]
def categorical_bincount(count, log_ps, n):
  """Given (un-norm.) log_probs, sample and count how many times for each index.

  Args:
    count: int, the amount of times to draw from the probabilities.
    log_ps: tensor shape [B x n], for each batch, n different numbers which can
      be unnormalized log probabilities that will get drawn from.
    n: should be the second entry in shape of log_ps.
  Returns:
    tensor shape [B x n], the counts of how many times each index was
      sampled.
  """
  # Force CPU since it's about 3x faster on CPU than GPU for current models.
  # TODO(ayzaan): tf.math.bincount is not compatible with TPU currently anyway.
  # Once it becomes compatible, we may have to remove/modify this use of
  # tf.device.
  with tf.device('CPU:0'):
    # samples shape: [B x n]
    samples = tf.multinomial(log_ps, num_samples=count)
    # return shape: [B x n]
    bincounts = tf.math.bincount(
        samples, minlength=n, maxlength=n, axis=-1)
    return bincounts
//  [python]
def categorical_bincount(count, log_ps, n):
  """Given (un-norm.) log_probs, sample and count how many times for each index.

  Args:
    count: int, the amount of times to draw from the probabilities.
    log_ps: tensor shape [B x n], for each batch, n different numbers which can
      be unnormalized log probabilities that will get drawn from.
    n: should be the second entry in shape of log_ps.
  Returns:
    tensor shape [B x n], the counts of how many times each index was
      sampled.
  """
  # Force CPU since it's about 3x faster on CPU than GPU for current models.
  # TODO(ayzaan): tf.math.bincount is not compatible with TPU currently anyway.
  # Once it becomes compatible, we may have to remove/modify this use of
  # tf.device.
  with tf.device('CPU:0'):
    # samples shape: [B x n]
    samples = tf.random.categorical(log_ps, num_samples=count, dtype=tf.int32)
    # return shape: [B x n]
    bincounts = tf.math.bincount(
        samples, minlength=n, maxlength=n, axis=-1)
    return bincounts
//  [python]
def get_policy(self):

    def policy_fn(observation, dtype=tf.int32):
      if tf.rank(observation) < 1:
        observation = [observation]
      distribution = tf.nn.softmax(tf.gather(self.policy, observation), axis=-1)
      batched = tf.rank(distribution) > 1
      if not batched:
        distributions = distribution[None, :]
      else:
        distributions = distribution

      batch_size = tf.shape(distributions)[0]

      actions = tf.random.categorical(
          tf.math.log(1e-8 + distributions), 1, dtype=dtype)
      actions = tf.squeeze(actions, -1)
      probs = tf.gather_nd(
          distributions,
          tf.stack([tf.range(batch_size, dtype=dtype), actions], -1))

      if not batched:
        action = actions[0]
        log_prob = tf.math.log(1e-8 + probs[0])
      else:
        action = actions
        log_prob = tf.math.log(1e-8 + probs)

      policy_info = {'distribution': distribution}
      return (tfp.distributions.Categorical(probs=distribution,
                                            dtype=dtype), policy_info)

    policy_info_spec = {
        'log_probability':
            specs.TensorSpec([], tf.float32),
        'distribution':
            specs.BoundedTensorSpec([self.num_actions],
                                    tf.float32,
                                    minimum=0.0,
                                    maximum=1.0)
    }
    return policy_fn, policy_info_spec
//  [python]
def filter_random_lighting(sequence_dir):
  sequence_name = tf.string_split([sequence_dir], '/').values[-1]
  lighting = tf.substr(sequence_name, 0, 6)
  return tf.not_equal(lighting, 'random')
//  [python]
def filter_random_lighting(sequence_dir):
  sequence_name = tf.string_split([sequence_dir], '/').values[-1]
  lighting = tf.strings.substr(sequence_name, 0, 6)
  return tf.not_equal(lighting, 'random')
//  [python]
def call(self, str_inp):
        # we're basically always going to let TFHub modules do space                                               
        # tokenization for us                                                                                      

        # blech, it's not really possible to actually define a Keras input w/a shape of ndim 1
        str_inp_squeezed = tensorflow.squeeze(str_inp, axis=1)

        # let's apply the max strlen to prevent OOM hopefully                                                  
        str_inp_cutoff = t.substr(str_inp_squeezed, 0, self.max_strlen)

        return self.embedder(str_inp_cutoff, as_dict=True)[self.output_key]
//  [python]
def call(self, str_inp):
        # we're basically always going to let TFHub modules do space                                               
        # tokenization for us                                                                                      

        # blech, it's not really possible to actually define a Keras input w/a shape of ndim 1
        str_inp_squeezed = tensorflow.squeeze(str_inp, axis=1)

        # let's apply the max strlen to prevent OOM hopefully                                                  
        str_inp_cutoff = tensorflow.strings.substr(str_inp_squeezed, 0, self.max_strlen)

        return self.embedder(str_inp_cutoff, as_dict=True)[self.output_key]
//  [python]
def _top_k_sample(logits, ignore_ids=None, num_samples=1, k=10):
    """
    Does top-k sampling. if ignore_ids is on, then we will zero out those logits.
    :param logits: [batch_size, vocab_size] tensor
    :param ignore_ids: [vocab_size] one-hot representation of the indices we'd like to ignore and never predict,
                        like padding maybe
    :param p: topp threshold to use, either a float or a [batch_size] vector
    :return: [batch_size, num_samples] samples
    # TODO FIGURE OUT HOW TO DO THIS ON TPUS. IT'S HELLA SLOW RIGHT NOW, DUE TO ARGSORT I THINK
    """
    with tf.compat.v1.variable_scope('top_p_sample'):
        batch_size, vocab_size = get_shape_list(logits, expected_rank=2)

        probs = tf.nn.softmax(logits if ignore_ids is None else logits - tf.cast(ignore_ids[None], tf.float32) * 1e10,
                              axis=-1)
        # [batch_size, vocab_perm]
        indices = tf.argsort(probs, direction='DESCENDING')

        # find the top pth index to cut off. careful we don't want to cutoff everything!
        # result will be [batch_size, vocab_perm]
        k_expanded = k if isinstance(k, int) else k[:, None]
        exclude_mask = tf.range(vocab_size)[None] >= k_expanded

        # OPTION A - sample in the sorted space, then unsort.
        logits_to_use = tf.batch_gather(logits, indices) - tf.cast(exclude_mask, tf.float32) * 1e10
        sample_perm = tf.multinomial(logits=logits_to_use, num_samples=num_samples)
        sample = tf.batch_gather(indices, sample_perm)

    return {
        'probs': probs,
        'sample': sample,
    }
//  [python]
def _top_k_sample(logits, ignore_ids=None, num_samples=1, k=10):
    """
    Does top-k sampling. if ignore_ids is on, then we will zero out those logits.
    :param logits: [batch_size, vocab_size] tensor
    :param ignore_ids: [vocab_size] one-hot representation of the indices we'd like to ignore and never predict,
                        like padding maybe
    :param p: topp threshold to use, either a float or a [batch_size] vector
    :return: [batch_size, num_samples] samples
    # TODO FIGURE OUT HOW TO DO THIS ON TPUS. IT'S HELLA SLOW RIGHT NOW, DUE TO ARGSORT I THINK
    """
    with tf.compat.v1.variable_scope('top_p_sample'):
        batch_size, vocab_size = get_shape_list(logits, expected_rank=2)

        probs = tf.nn.softmax(logits if ignore_ids is None else logits - tf.cast(ignore_ids[None], tf.float32) * 1e10,
                              axis=-1)
        # [batch_size, vocab_perm]
        indices = tf.argsort(probs, direction='DESCENDING')

        # find the top pth index to cut off. careful we don't want to cutoff everything!
        # result will be [batch_size, vocab_perm]
        k_expanded = k if isinstance(k, int) else k[:, None]
        exclude_mask = tf.range(vocab_size)[None] >= k_expanded

        # OPTION A - sample in the sorted space, then unsort.
        logits_to_use = tf.batch_gather(logits, indices) - tf.cast(exclude_mask, tf.float32) * 1e10
        sample_perm = tf.random.categorical(logits=logits_to_use, num_samples=num_samples)
        sample = tf.batch_gather(indices, sample_perm)

    return {
        'probs': probs,
        'sample': sample,
    }
//  [python]
def sample(model, starting_str, 
           len_generated_text=500, 
           max_input_length=40,
           scale_factor=1.0):
    encoded_input = [char2int[s] for s in starting_str]
    encoded_input = tf.reshape(encoded_input, (1, -1))

    generated_str = starting_str

    model.reset_states()
    for i in range(len_generated_text):
        logits = model(encoded_input)
        logits = tf.squeeze(logits, 0)

        scaled_logits = logits * scale_factor
        new_char_indx = tf.multinomial(
            scaled_logits, num_samples=1)
        
        new_char_indx = tf.squeeze(new_char_indx)[-1].numpy()    

        generated_str += str(char_array[new_char_indx])
        
        new_char_indx = tf.expand_dims([new_char_indx], 0)
        encoded_input = tf.concat(
            [encoded_input, new_char_indx],
            axis=1)
        encoded_input = encoded_input[:, -max_input_length:]

    return generated_str
//  [python]
def sample(model, starting_str, 
           len_generated_text=500, 
           max_input_length=40,
           scale_factor=1.0):
    encoded_input = [char2int[s] for s in starting_str]
    encoded_input = tf.reshape(encoded_input, (1, -1))

    generated_str = starting_str

    model.reset_states()
    for i in range(len_generated_text):
        logits = model(encoded_input)
        logits = tf.squeeze(logits, 0)

        scaled_logits = logits * scale_factor
        new_char_indx = tf.random.categorical(
            scaled_logits, num_samples=1)
        
        new_char_indx = tf.squeeze(new_char_indx)[-1].numpy()    

        generated_str += str(char_array[new_char_indx])
        
        new_char_indx = tf.expand_dims([new_char_indx], 0)
        encoded_input = tf.concat(
            [encoded_input, new_char_indx],
            axis=1)
        encoded_input = encoded_input[:, -max_input_length:]

    return generated_str
//  [python]
autoenc_model_dir = "/home/david/Downloads/tfmodels/deepaec/"
frozen_model = os.path.join(autoenc_model_dir, "autoenc_frozen_model.pb")

print("load model: ", autoenc_model_dir)
with tf.Graph().as_default() as graph:
    with tf.Session(graph=graph) as sess:
        tf.saved_model.loader.load(sess, [tag_constants.SERVING], autoenc_model_dir)

        graphdef = graph.as_graph_def()
        print("extract autoenc subgraph")
        subgraphdef = tf.graph_util.extract_sub_graph(graphdef, ["output256"])
        print("convert variables to constants")
        subgraphdef2 = tf.graph_util.convert_variables_to_constants(sess, subgraphdef, ['output256'])

        print("write frozen graph")
        with tf.gfile.GFile(frozen_model, "wb") as f:
            f.write(subgraphdef2.SerializeToString())
//  [python]
autoenc_model_dir = "/home/david/Downloads/tfmodels/deepaec/"
frozen_model = os.path.join(autoenc_model_dir, "autoenc_frozen_model.pb")

print("load model: ", autoenc_model_dir)
with tf.Graph().as_default() as graph:
    with tf.compat.v1.Session(graph=graph) as sess:
        tf.saved_model.load(sess, autoenc_model_dir)

        graphdef = graph.as_graph_def()
        print("extract autoenc subgraph")
        subgraphdef = tf.graph_util.extract_sub_graph(graphdef, ["output256"])
        print("convert variables to constants")
        subgraphdef2 = tf.graph_util.convert_variables_to_constants(sess, subgraphdef, ['output256'])

        print("write frozen graph")
        with tf.gfile.GFile(frozen_model, "wb") as f:
            f.write(subgraphdef2.SerializeToString())
//  [python]
def getSess():
    flags = tf.flags
    FLAGS = flags.FLAGS

    flags.DEFINE_string(
        "export_dir", '~/Downloads/bert_small/',
        "The dir where the exported model has been written.")

    tokenizer = tokenization.FullTokenizer(vocab_file='~/Downloads/bert_small/vocab.txt', do_lower_case=True)
    sess = tf.InteractiveSession()
    tf.saved_model.loader.load(sess, [tag_constants.SERVING], FLAGS.export_dir)
    return sess, tokenizer
//  [python]
def getSess():
    flags = tf.flags
    FLAGS = flags.FLAGS

    flags.DEFINE_string(
        "export_dir", '~/Downloads/bert_small/',
        "The dir where the exported model has been written.")

    tokenizer = tokenization.FullTokenizer(vocab_file='~/Downloads/bert_small/vocab.txt', do_lower_case=True)
    sess = tf.InteractiveSession()
    tf.saved_model.load(sess, FLAGS.export_dir)
    return sess, tokenizer
//  [python]
def fid_to_image_attribute(fid, pid, attr_lookuptable, image_root, image_size):
    """ Loads and resizes an image given by FID. Pass-through the PID. """
    # Since there is no symbolic path.join, we just add a '/' to be sure.
    # if not hasattr(fid_to_image_attribute, "_label"):


    image_encoded = tf.io.read_file(tf.reduce_join([image_root, '/', fid]))
    #image_encoded = tf.read_file(tf.reduce_join([image_root, '/', fid]))

    # tf.image.decode_image doesn't set the shape, not even the dimensionality,
    # because it potentially loads animated .gif files. Instead, we use either
    # decode_jpeg or decode_png, each of which can decode both.
    # Sounds ridiculous, but is true:
    # https://github.com/tensorflow/tensorflow/issues/9356#issuecomment-309144064
    image_decoded = tf_v2.image.decode_jpeg(image_encoded, channels=3)
    image_resized = tf_v2.image.resize(image_decoded, image_size)

    id = tf.string_split(fid, delimiter='-')[0]

    lookuptable = attr_lookuptable[0]
    relative_id = lookuptable.lookup(id)

    attribute_label = []
    for i in range(1,28):
        attribute_label.append(attr_lookuptable[i].lookup(relative_id))

    return image_resized, attribute_label, fid, pid
//  [python]
def fid_to_image_attribute(fid, pid, attr_lookuptable, image_root, image_size):
    """ Loads and resizes an image given by FID. Pass-through the PID. """
    # Since there is no symbolic path.join, we just add a '/' to be sure.
    # if not hasattr(fid_to_image_attribute, "_label"):


    image_encoded = tf.io.read_file(tf.reduce_join([image_root, '/', fid]))
    #image_encoded = tf.read_file(tf.reduce_join([image_root, '/', fid]))

    # tf.image.decode_image doesn't set the shape, not even the dimensionality,
    # because it potentially loads animated .gif files. Instead, we use either
    # decode_jpeg or decode_png, each of which can decode both.
    # Sounds ridiculous, but is true:
    # https://github.com/tensorflow/tensorflow/issues/9356#issuecomment-309144064
    image_decoded = tf_v2.image.decode_jpeg(image_encoded, channels=3)
    image_resized = tf_v2.image.resize(image_decoded, image_size)

    id = tf_v2.strings.substr(fid, 19, 4)

    lookuptable = attr_lookuptable[0]
    relative_id = lookuptable.lookup(id)

    attribute_label = []
    for i in range(1,28):
        attribute_label.append(attr_lookuptable[i].lookup(relative_id))

    return image_resized, attribute_label, fid, pid
//  [python]
def predict(self, **data):
        '''
        使用模型
        :param path: 模型所在的路径
        :param name: 模型的名字
        :param data: 模型的输入参数
        :return:
        '''
        with tf.Session() as session:
            tf.saved_model.loader.load(session, [tag_constants.SERVING], os.path.join(MODEL_PATH, TENSORFLOW_MODEL_DIR))
            input_x = session.graph.get_tensor_by_name(self.get_tensor_name('input_x'))
            y_conv = session.graph.get_tensor_by_name(self.get_tensor_name('y_conv'))
            keep_prob = session.graph.get_tensor_by_name(self.get_tensor_name('keep_prob'))
            x_data = self.data.predict_data(**data)
            predict = session.run(y_conv, feed_dict={input_x: x_data, keep_prob: 1.0})
            return self.data.to_categorys(predict)
//  [python]
def predict(self, **data):
        '''
        使用模型
        :param path: 模型所在的路径
        :param name: 模型的名字
        :param data: 模型的输入参数
        :return:
        '''
        model = tf.saved_model.load(os.path.join(MODEL_PATH, TENSORFLOW_MODEL_DIR))
        input_x = model.signatures['serving_default'].inputs[0]
        y_conv = model.signatures['serving_default'].outputs[0]
        keep_prob = model.signatures['serving_default'].inputs[1]
        x_data = self.data.predict_data(**data)
        predict = model.signatures['serving_default'](input_x=x_data, keep_prob=1.0)
        return self.data.to_categorys(predict['output_0'])
//  [python]
def _load_corpus(self, corpus_dir):
        for fd in range(2, -1, -1):
            file_list = []
            if fd == 0:
                file_dir = os.path.join(corpus_dir, AUG0_FOLDER)
            elif fd == 1:
                file_dir = os.path.join(corpus_dir, AUG1_FOLDER)
            else:
                file_dir = os.path.join(corpus_dir, AUG2_FOLDER)

            for data_file in sorted(os.listdir(file_dir)):
                full_path_name = os.path.join(file_dir, data_file)
                if os.path.isfile(full_path_name) and data_file.lower().endswith('.txt'):
                    file_list.append(full_path_name)

            assert len(file_list) > 0
            dataset = tf.data.TextLineDataset(file_list)

            src_dataset = dataset.filter(lambda line:
                                         tf.logical_and(tf.size(line) > 0,
                                                        tf.equal(tf.substr(line, 0, 2), tf.constant('Q:'))))
            src_dataset = src_dataset.map(lambda line:
                                          tf.substr(line, 2, MAX_LEN)).prefetch(4096)
            tgt_dataset = dataset.filter(lambda line:
                                         tf.logical_and(tf.size(line) > 0,
                                                        tf.equal(tf.substr(line, 0, 2), tf.constant('A:'))))
            tgt_dataset = tgt_dataset.map(lambda line:
                                          tf.substr(line, 2, MAX_LEN)).prefetch(4096)

            src_tgt_dataset = tf.data.Dataset.zip((src_dataset, tgt_dataset))
            if fd == 1:
                src_tgt_dataset = src_tgt_dataset.repeat(self.hparams.aug1_repeat_times)
            elif fd == 2:
                src_tgt_dataset = src_tgt_dataset.repeat(self.hparams.aug2_repeat_times)

            if self.text_set is None:
                self.text_set = src_tgt_dataset
            else:
                self.text_set = self.text_set.concatenate(src_tgt_dataset)
//  [python]
def _load_corpus(self, corpus_dir):
        for fd in range(2, -1, -1):
            file_list = []
            if fd == 0:
                file_dir = os.path.join(corpus_dir, AUG0_FOLDER)
            elif fd == 1:
                file_dir = os.path.join(corpus_dir, AUG1_FOLDER)
            else:
                file_dir = os.path.join(corpus_dir, AUG2_FOLDER)

            for data_file in sorted(os.listdir(file_dir)):
                full_path_name = os.path.join(file_dir, data_file)
                if os.path.isfile(full_path_name) and data_file.lower().endswith('.txt'):
                    file_list.append(full_path_name)

            assert len(file_list) > 0
            dataset = tf.data.TextLineDataset(file_list)

            src_dataset = dataset.filter(lambda line:
                                         tf.logical_and(tf.size(line) > 0,
                                                        tf.equal(tf.strings.substr(line, 0, 2), tf.constant('Q:'))))
            src_dataset = src_dataset.map(lambda line:
                                          tf.strings.substr(line, 2, MAX_LEN)).prefetch(4096)
            tgt_dataset = dataset.filter(lambda line:
                                         tf.logical_and(tf.size(line) > 0,
                                                        tf.equal(tf.strings.substr(line, 0, 2), tf.constant('A:'))))
            tgt_dataset = tgt_dataset.map(lambda line:
                                          tf.strings.substr(line, 2, MAX_LEN)).prefetch(4096)

            src_tgt_dataset = tf.data.Dataset.zip((src_dataset, tgt_dataset))
            if fd == 1:
                src_tgt_dataset = src_tgt_dataset.repeat(self.hparams.aug1_repeat_times)
            elif fd == 2:
                src_tgt_dataset = src_tgt_dataset.repeat(self.hparams.aug2_repeat_times)

            if self.text_set is None:
                self.text_set = src_tgt_dataset
            else:
                self.text_set = self.text_set.concatenate(src_tgt_dataset)
//  [python]
def predict_all(self, datas):
        with tf.Session() as session:
            tf.saved_model.loader.load(session, [tag_constants.SERVING], os.path.join(MODEL_PATH, TENSORFLOW_MODEL_DIR))
            input_x = session.graph.get_tensor_by_name(self.get_tensor_name('input_x'))
            y_conv = session.graph.get_tensor_by_name(self.get_tensor_name('y_conv'))
            keep_prob = session.graph.get_tensor_by_name(self.get_tensor_name('keep_prob'))
            ratings = []
            for data in datas:
                x_data = self.data.predict_data(**data)
                predict = session.run(y_conv, feed_dict={input_x: x_data, keep_prob: 1.0})
                predict = self.data.to_categorys(predict)
                ratings.append( predict)
        return ratings
//  [python]
def predict_all(self, datas):
        with tf.compat.v1.Session() as session:
            tf.saved_model.load(session, os.path.join(MODEL_PATH, TENSORFLOW_MODEL_DIR))
            input_x = session.graph.get_tensor_by_name(self.get_tensor_name('input_x'))
            y_conv = session.graph.get_tensor_by_name(self.get_tensor_name('y_conv'))
            keep_prob = session.graph.get_tensor_by_name(self.get_tensor_name('keep_prob'))
            ratings = []
            for data in datas:
                x_data = self.data.predict_data(**data)
                predict = session.run(y_conv, feed_dict={input_x: x_data, keep_prob: 1.0})
                predict = self.data.to_categorys(predict)
                ratings.append( predict)
        return ratings
//  [python]
def _creation_date_string_to_integer(dates):
  """Converts ISO date string to integer that can be sorted.

  Returned integers retain the property that sorting the integers results in
  dates sorted in chronological order, so these integers can be used to sort
  examples by date. Ignores fractional seconds if provided.
  Assumes standard time offset.

  For example:
    2009-06-15T13:45:30 -> 20090615134530
    2009-06-15T13:45:30.345Z -> 20090615134530

  Args:
    dates: A tf.string tensor of dates in simplified ISO 8601 format. The data
      produced by `tff.simulation.datasets.stackoverflow.load_data` conforms
      to this format.

  Returns:
    A tf.int64 tensor of integers representing dates.
  """
  year = tf.strings.to_number(
      tf.strings.slice(dates, [0, 0], [4, 4]), out_type=tf.int64)
  month = tf.strings.to_number(
      tf.strings.slice(dates, [0, 5], [4, 7]), out_type=tf.int64)
  day = tf.strings.to_number(
      tf.strings.slice(dates, [0, 8], [4, 10]), out_type=tf.int64)
  hour = tf.strings.to_number(
      tf.strings.slice(dates, [0, 11], [4, 13]), out_type=tf.int64)
  minute = tf.strings.to_number(
      tf.strings.slice(dates, [0, 14], [4, 16]), out_type=tf.int64)
  second = tf.strings.to_number(
      tf.strings.slice(dates, [0, 17], [4, 19]), out_type=tf.int64)

  timestamp = 0
  timestamp = (timestamp + year) * 100
  timestamp = (timestamp + month) * 100
  timestamp = (timestamp + day) * 100
  timestamp = (timestamp + hour) * 100
  timestamp = (timestamp + minute) * 100
  timestamp = timestamp + second
  return timestamp
//  [python]
def _creation_date_string_to_integer(dates):
  """Converts ISO date string to integer that can be sorted.

  Returned integers retain the property that sorting the integers results in
  dates sorted in chronological order, so these integers can be used to sort
  examples by date. Ignores fractional seconds if provided.
  Assumes standard time offset.

  For example:
    2009-06-15T13:45:30 -> 20090615134530
    2009-06-15T13:45:30.345Z -> 20090615134530

  Args:
    dates: A tf.string tensor of dates in simplified ISO 8601 format. The data
      produced by `tff.simulation.datasets.stackoverflow.load_data` conforms
      to this format.

  Returns:
    A tf.int64 tensor of integers representing dates.
  """
  year = tf.strings.to_number(
      tf.strings.substr(dates, 0, 4), out_type=tf.int64)
  month = tf.strings.to_number(
      tf.strings.substr(dates, 5, 2), out_type=tf.int64)
  day = tf.strings.to_number(
      tf.strings.substr(dates, 8, 2), out_type=tf.int64)
  hour = tf.strings.to_number(
      tf.strings.substr(dates, 11, 2), out_type=tf.int64)
  minute = tf.strings.to_number(
      tf.strings.substr(dates, 14, 2), out_type=tf.int64)
  second = tf.strings.to_number(
      tf.strings.substr(dates, 17, 2), out_type=tf.int64)

  timestamp = 0
  timestamp = (timestamp + year) * 100
  timestamp = (timestamp + month) * 100
  timestamp = (timestamp + day) * 100
  timestamp = (timestamp + hour) * 100
  timestamp = (timestamp + minute) * 100
  timestamp = timestamp + second
  return timestamp