jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def crop_and_resize(image, boxes, box_ind, crop_size, pad_border=False):
  # image feature [1,C,FS,FS] # for mask gt [N_FG, 1, H, W]
  # boxes [N,4]
  # box_ind [N] all zero?

  if pad_border:
    image = tf.pad(image, [[0, 0], [0, 0], [1, 1], [1, 1]], mode="SYMMETRIC")
    boxes = boxes + 1

  # return [N,C,crop_size,crop_size]
  def transform_fpcoor_for_tf(boxes, image_shape, crop_shape):
    """
    The way tf.image.crop_and_resize works (with normalized box):
    Initial point (the value of output[0]): x0_box * (W_img - 1)
    Spacing: w_box * (W_img - 1) / (W_crop - 1)
    Use the above grid to bilinear sample.

    However, what we want is (with fpcoor box):
    Spacing: w_box / W_crop
    Initial point: x0_box + spacing/2 - 0.5
    (-0.5 because bilinear sample assumes floating point coordinate (0.0, 0.0)
    is the same as pixel value (0, 0))

    This function transform fpcoor boxes to a format to be used by
    tf.image.crop_and_resize

    Returns:
      y1x1y2x2
    """
    x0, y0, x1, y1 = tf.split(boxes, 4, axis=1)

    spacing_w = (x1 - x0) / tf.cast(crop_shape[1], tf.float32)
    spacing_h = (y1 - y0) / tf.cast(crop_shape[0], tf.float32)

    nx0 = (x0 + spacing_w / 2 - 0.5) / tf.cast(image_shape[1] - 1, tf.float32)
    ny0 = (y0 + spacing_h / 2 - 0.5) / tf.cast(image_shape[0] - 1, tf.float32)

    nw = spacing_w * tf.cast(crop_shape[1] - 1, tf.float32) / \
        tf.cast(image_shape[1] - 1, tf.float32)
    nh = spacing_h * tf.cast(crop_shape[0] - 1, tf.float32) / \
        tf.cast(image_shape[0] - 1, tf.float32)

    return tf.concat([ny0, nx0, ny0 + nh, nx0 + nw], axis=1)

  image_shape = tf.shape(image)[2:]
  boxes = transform_fpcoor_for_tf(boxes, image_shape, [crop_size, crop_size])
  image = tf.transpose(image, [0, 2, 3, 1])   # 1hwc
  ret = tf.image.crop_and_resize(
      image, boxes, box_ind,
      crop_size=[crop_size, crop_size])
  ret = tf.transpose(ret, [0, 3, 1, 2])   # Ncss
  return ret
//  [python]
def _shard_features(self, features):  # pylint: disable=missing-docstring
        sharded_features = {}
        for k, v in sorted(six.iteritems(features)):
            v = tf.convert_to_tensor(v)
            v_shape = common_layers.shape_list(v)
            if not v_shape:
                v = tf.expand_dims(v, axis = -1)
                v_shape = [1]
            if v_shape == [1]:
                v = tf.tile(v, tf.to_int32([self._num_datashards]))
            sharded_features[k] = self._data_parallelism(
                tf.identity, tf.split(v, self._num_datashards, 0)
            )
        return sharded_features
//  [python]
def _shard_features(self, features):  # pylint: disable=missing-docstring
        sharded_features = {}
        for k, v in sorted(six.iteritems(features)):
            v = tf.convert_to_tensor(v)
            v_shape = common_layers.shape_list(v)
            if not v_shape:
                v = tf.expand_dims(v, axis = -1)
                v_shape = [1]
            if v_shape == [1]:
                v = tf.tile(v, tf.cast([self._num_datashards], tf.int32))
            sharded_features[k] = self._data_parallelism(
                tf.identity, tf.split(v, self._num_datashards, 0)
            )
        return sharded_features
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training,
                                drop_remainder):
    """Creates an `input_fn` closure to be passed to TPUEstimator."""

    name_to_features = {
        "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
        "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
        "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
        "entity_indices": tf.FixedLenFeature([FLAGS.entity_num, FLAGS.max_num_entity_indices], tf.int64),
        "in_neighbors_indices": tf.FixedLenFeature([seq_length, FLAGS.max_num_neighbors], tf.int64),
        "entity_indices_mask": tf.FixedLenFeature([FLAGS.entity_num, FLAGS.max_num_entity_indices], tf.int64),
        "in_neighbors_indices_mask": tf.FixedLenFeature([seq_length, FLAGS.max_num_neighbors], tf.int64),
        "label_ids": tf.FixedLenFeature([], tf.int64),
        "is_real_example": tf.FixedLenFeature([], tf.int64),
    }

    def _decode_record(record, name_to_features):
        """Decodes a record to a TensorFlow example."""
        example = tf.parse_single_example(record, name_to_features)

        # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
        # So cast all int64 to int32.
        for name in list(example.keys()):
            t = example[name]
            if t.dtype == tf.int64:
                t = tf.to_int32(t)
            example[name] = t

        return example

    def input_fn(params):
        """The actual input function."""
        batch_size = params["batch_size"]

        # For training, we want a lot of parallel reading and shuffling.
        # For eval, we want no shuffling and parallel reading doesn't matter.
        d = tf.data.TFRecordDataset(input_file)
        if is_training:
            d = d.repeat()
            d = d.shuffle(buffer_size=100)

        d = d.apply(
            tf.data.experimental.map_and_batch(
                lambda record: _decode_record(record, name_to_features),
                batch_size=batch_size,
                drop_remainder=drop_remainder))

        return d

    return input_fn
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training,
                                drop_remainder):
    """Creates an `input_fn` closure to be passed to TPUEstimator."""

    name_to_features = {
        "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
        "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
        "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
        "entity_indices": tf.FixedLenFeature([FLAGS.entity_num, FLAGS.max_num_entity_indices], tf.int64),
        "in_neighbors_indices": tf.FixedLenFeature([seq_length, FLAGS.max_num_neighbors], tf.int64),
        "entity_indices_mask": tf.FixedLenFeature([FLAGS.entity_num, FLAGS.max_num_entity_indices], tf.int64),
        "in_neighbors_indices_mask": tf.FixedLenFeature([seq_length, FLAGS.max_num_neighbors], tf.int64),
        "label_ids": tf.FixedLenFeature([], tf.int64),
        "is_real_example": tf.FixedLenFeature([], tf.int64),
    }

    def _decode_record(record, name_to_features):
        """Decodes a record to a TensorFlow example."""
        example = tf.parse_single_example(record, name_to_features)

        # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
        # So cast all int64 to int32.
        for name in list(example.keys()):
            t = example[name]
            if t.dtype == tf.int64:
                t = tf.cast(t, tf.int32)
            example[name] = t

        return example

    def input_fn(params):
        """The actual input function."""
        batch_size = params["batch_size"]

        # For training, we want a lot of parallel reading and shuffling.
        # For eval, we want no shuffling and parallel reading doesn't matter.
        d = tf.data.TFRecordDataset(input_file)
        if is_training:
            d = d.repeat()
            d = d.shuffle(buffer_size=100)

        d = d.apply(
            tf.data.experimental.map_and_batch(
                lambda record: _decode_record(record, name_to_features),
                batch_size=batch_size,
                drop_remainder=drop_remainder))

        return d

    return input_fn
//  [python]
def length_penalty(self, sequence_lengths, penalty_factor):
        penalty_factor = tf.to_float(penalty_factor)
        return tf.div((tf.to_float(5.0) + tf.to_float(sequence_lengths))**penalty_factor, (tf.to_float(1.0 + 5.0))**penalty_factor)
//  [python]
def length_penalty(self, sequence_lengths, penalty_factor):
        penalty_factor = tf.cast(penalty_factor, tf.float32)
        return tf.div((tf.cast(5.0, tf.float32) + tf.cast(sequence_lengths, tf.float32))**penalty_factor, (tf.cast(1.0, tf.float32) + tf.cast(5.0, tf.float32))**penalty_factor)
//  [python]
def summarize_features(features, num_shards = 1):
    """Generate summaries for features."""
    if not common_layers.should_generate_summaries():
        return

    with tf.name_scope('input_stats'):
        for (k, v) in sorted(six.iteritems(features)):
            if (
                isinstance(v, tf.Tensor)
                and (v.get_shape().ndims > 1)
                and (v.dtype != tf.string)
            ):
                tf.summary.scalar('%s_batch' % k, tf.shape(v)[0] // num_shards)
                tf.summary.scalar('%s_length' % k, tf.shape(v)[1])
                nonpadding = tf.to_float(tf.not_equal(v, 0))
                nonpadding_tokens = tf.reduce_sum(nonpadding)
                tf.summary.scalar('%s_nonpadding_tokens' % k, nonpadding_tokens)
                tf.summary.scalar(
                    '%s_nonpadding_fraction' % k, tf.reduce_mean(nonpadding)
                )
//  [python]
def summarize_features(features, num_shards = 1):
    """Generate summaries for features."""
    if not common_layers.should_generate_summaries():
        return

    with tf.name_scope('input_stats'):
        for (k, v) in sorted(six.iteritems(features)):
            if (
                isinstance(v, tf.Tensor)
                and (v.get_shape().ndims > 1)
                and (v.dtype != tf.string)
            ):
                tf.summary.scalar('%s_batch' % k, tf.shape(v)[0] // num_shards)
                tf.summary.scalar('%s_length' % k, tf.shape(v)[1])
                nonpadding = tf.cast(tf.not_equal(v, 0), tf.float32)
                nonpadding_tokens = tf.reduce_sum(nonpadding)
                tf.summary.scalar('%s_nonpadding_tokens' % k, nonpadding_tokens)
                tf.summary.scalar(
                    '%s_nonpadding_fraction' % k, tf.reduce_mean(nonpadding)
                )
//  [python]
def code(self, inputs):
        batch_size, seqslen = tf.shape(inputs)[0], tf.shape(inputs)[1]
        inputs = tf.to_float(tf.reshape(inputs, [batch_size, seqslen, -1]))
        inputs = tf.tile(tf.expand_dims(inputs, -1), [1, 1, 1, self._num_units])

        # Harmonic
        outs = inputs * self.basis_freq
        outs = tf.nn.bias_add(outs, self.phase)
        outs = tf.cos(outs)
        return outs
//  [python]
def code(self, inputs):
        batch_size, seqslen = tf.shape(inputs)[0], tf.shape(inputs)[1]
        inputs = tf.cast(tf.reshape(inputs, [batch_size, seqslen, -1]), tf.float32)
        inputs = tf.tile(tf.expand_dims(inputs, -1), [1, 1, 1, self._num_units])

        # Harmonic
        outs = inputs * self.basis_freq
        outs = tf.nn.bias_add(outs, self.phase)
        outs = tf.cos(outs)
        return outs
//  [python]
def code(self, inputs):
        shape_list = inputs.shape.as_list()
        assert len(shape_list) == 2, "the tensor rank should be 2."

        x = tf.expand_dims(tf.to_float(inputs), axis=-1)
        x = tf.tile(x, [1, 1, self.num_units // 2]) / self.scale

        code_even = tf.sin(x)
        code_odd = tf.cos(x)

        code = tf.stack([code_even, code_odd], axis=-1)
        code = tf.reshape(code, [-1, shape_list[1], self.num_units])
        return code
//  [python]
def code(self, inputs):
        shape_list = inputs.shape.as_list()
        assert len(shape_list) == 2, "the tensor rank should be 2."

        x = tf.expand_dims(tf.cast(inputs, tf.float32), axis=-1)
        x = tf.tile(x, [1, 1, self.num_units // 2]) / self.scale

        code_even = tf.sin(x)
        code_odd = tf.cos(x)

        code = tf.stack([code_even, code_odd], axis=-1)
        code = tf.reshape(code, [-1, shape_list[1], self.num_units])
        return code
//  [python]
def compute_attention_weight(self, state, hidden_states, source_mask):
        encoder_hidden_hid = self.get_config('encoder', 'cell')['num_units']
        decoder_hidden_hid = self.config['cell']['num_units']

        with tf.variable_scope('attention', reuse=tf.AUTO_REUSE) as scope:
            W_d_hh = tf.get_variable(name='W_d_hh', shape=[encoder_hidden_hid*2, decoder_hidden_hid], initializer=tf.contrib.layers.xavier_initializer(), dtype=tf.float32)
            W_e_hh = tf.get_variable(name='W_e_hh', shape=[decoder_hidden_hid, decoder_hidden_hid], initializer=tf.contrib.layers.xavier_initializer(), dtype=tf.float32)
            b_a = tf.get_variable(name='b_a', shape=[decoder_hidden_hid], initializer=tf.zeros_initializer(), dtype=tf.float32)
            v = tf.get_variable(name='v', shape=[decoder_hidden_hid], initializer=tf.contrib.layers.xavier_initializer(), dtype=tf.float32)

        r_1 = tf.einsum('ijk,kl->ijl', hidden_states, W_e_hh)
        r_2 = tf.expand_dims(tf.matmul(state, W_d_hh), axis=1)
        a = tf.add(r_1, r_2)
        t = tf.tanh(tf.add(a, b_a))
        e_ij = tf.reduce_sum(tf.multiply(v, t), axis=-1)

        mask_value = tf.log(tf.to_float(0.0))
        e_ij_mask = mask_value * tf.ones_like(e_ij)

        e_ij = tf.where(source_mask, e_ij, e_ij_mask) # replace the attention weight on <PAD> token with -inf, then after softmax it can be zero

        a_ij = tf.expand_dims(tf.nn.softmax(e_ij, axis=1), axis=-1)
        return a_ij
//  [python]
def compute_attention_weight(self, state, hidden_states, source_mask):
        encoder_hidden_hid = self.get_config('encoder', 'cell')['num_units']
        decoder_hidden_hid = self.config['cell']['num_units']

        with tf.variable_scope('attention', reuse=tf.AUTO_REUSE) as scope:
            W_d_hh = tf.get_variable(name='W_d_hh', shape=[encoder_hidden_hid*2, decoder_hidden_hid], initializer=tf.contrib.layers.xavier_initializer(), dtype=tf.float32)
            W_e_hh = tf.get_variable(name='W_e_hh', shape=[decoder_hidden_hid, decoder_hidden_hid], initializer=tf.contrib.layers.xavier_initializer(), dtype=tf.float32)
            b_a = tf.get_variable(name='b_a', shape=[decoder_hidden_hid], initializer=tf.zeros_initializer(), dtype=tf.float32)
            v = tf.get_variable(name='v', shape=[decoder_hidden_hid], initializer=tf.contrib.layers.xavier_initializer(), dtype=tf.float32)

        r_1 = tf.einsum('ijk,kl->ijl', hidden_states, W_e_hh)
        r_2 = tf.expand_dims(tf.matmul(state, W_d_hh), axis=1)
        a = tf.add(r_1, r_2)
        t = tf.tanh(tf.add(a, b_a))
        e_ij = tf.reduce_sum(tf.multiply(v, t), axis=-1)

        mask_value = tf.log(tf.cast(0.0, dtype=tf.float32))
        e_ij_mask = mask_value * tf.ones_like(e_ij)

        e_ij = tf.where(source_mask, e_ij, e_ij_mask) # replace the attention weight on <PAD> token with -inf, then after softmax it can be zero

        a_ij = tf.expand_dims(tf.nn.softmax(e_ij, axis=1), axis=-1)
        return a_ij
//  [python]
def compute_mean(dataset):
  def _helper(aggregate, x):
    total, n = aggregate
    return total + x, n + 1

  total, n = tfds.as_numpy(dataset.reduce((0., 0), _helper))
  return tf.to_float(total / n)
//  [python]
def compute_mean(dataset):
  def _helper(aggregate, x):
    total, n = aggregate
    return total + x, n + 1

  total, n = tfds.as_numpy(dataset.reduce((0., 0), _helper))
  return tf.cast(total / n, tf.float32)
//  [python]
def get_mnist(split="train", shuffle_files=False):
  """Get FashionMNIST dataset."""
  split_map = {
      "train": "train",
      "valid": "validation",
      "test": "test",
  }
  datasets = dict(
      zip(["train", "validation", "test"],
          tfds.load(
              "mnist:3.*.*",
              split=["train[:50000]", "train[50000:]", "test"],
              shuffle_files=shuffle_files,
              data_dir=FLAGS.data_dir)))
  preprocess = lambda x: tf.to_float(x["image"])
  train_mean = compute_mean(datasets[split_map["train"]].map(preprocess))
  return datasets[split_map[split]].map(preprocess), train_mean
//  [python]
def get_mnist(split="train", shuffle_files=False):
  """Get FashionMNIST dataset."""
  split_map = {
      "train": "train",
      "valid": "validation",
      "test": "test",
  }
  datasets = dict(
      zip(["train", "validation", "test"],
          tfds.load(
              "mnist:3.*.*",
              split=["train[:50000]", "train[50000:]", "test"],
              shuffle_files=shuffle_files,
              data_dir=FLAGS.data_dir)))
  preprocess = lambda x: tf.cast(x["image"], tf.float32)
  train_mean = compute_mean(datasets[split_map["train"]].map(preprocess))
  return datasets[split_map[split]].map(preprocess), train_mean
//  [python]
def get_fashion_mnist(split="train", shuffle_files=False):
  """Get FashionMNIST dataset."""
  split_map = {
      "train": "train",
      "valid": "train",  # No validation set, so reuse train.
      "test": "test",
  }
  datasets = tfds.load("fashion_mnist",
                       shuffle_files=shuffle_files,
                       data_dir=FLAGS.data_dir)
  preprocess = lambda x: tf.to_float(x["image"])
  train_mean = compute_mean(datasets[split_map["train"]].map(preprocess))
  return datasets[split_map[split]].map(preprocess), train_mean
//  [python]
def get_fashion_mnist(split="train", shuffle_files=False):
  """Get FashionMNIST dataset."""
  split_map = {
      "train": "train",
      "valid": "train",  # No validation set, so reuse train.
      "test": "test",
  }
  datasets = tfds.load("fashion_mnist",
                       shuffle_files=shuffle_files,
                       data_dir=FLAGS.data_dir)
  preprocess = lambda x: tf.cast(x["image"], tf.float32)
  train_mean = compute_mean(datasets[split_map["train"]].map(preprocess))
  return datasets[split_map[split]].map(preprocess), train_mean
//  [python]
def get_celeba(split="train", shuffle_files=False):
  """Get CelebA dataset."""
  split_map = {
      "train": "train",
      "valid": "validation",
      "test": "test",
  }
  datasets = tfds.load("celeb_a:2.*.*",
                       shuffle_files=shuffle_files,
                       data_dir=FLAGS.data_dir)
  def preprocess(sample, crop_width=80, image_size=CELEBA_IMAGE_SIZE):
    """Output images are in [0, 255]."""
    image_shape = sample["image"].shape
    crop_slices = [
        slice(w // 2 - crop_width, w // 2 + crop_width) for w in image_shape[:2]
    ] + [slice(None)]
    image_cropped = sample["image"][crop_slices]
    image_resized = tf.image.resize_images(image_cropped, [image_size] * 2)
    x = tf.to_float(image_resized)
    return x
  train_mean = compute_mean(datasets[split_map["train"]].map(preprocess))
  return datasets[split_map[split]].map(preprocess), train_mean
//  [python]
def get_celeba(split="train", shuffle_files=False):
  """Get CelebA dataset."""
  split_map = {
      "train": "train",
      "valid": "validation",
      "test": "test",
  }
  datasets = tfds.load("celeb_a:2.*.*",
                       shuffle_files=shuffle_files,
                       data_dir=FLAGS.data_dir)
  def preprocess(sample, crop_width=80, image_size=CELEBA_IMAGE_SIZE):
    """Output images are in [0, 255]."""
    image_shape = sample["image"].shape
    crop_slices = [
        slice(w // 2 - crop_width, w // 2 + crop_width) for w in image_shape[:2]
    ] + [slice(None)]
    image_cropped = sample["image"][crop_slices]
    image_resized = tf.image.resize_images(image_cropped, [image_size] * 2)
    x = tf.cast(image_resized, tf.float32)
    return x
  train_mean = compute_mean(datasets[split_map["train"]].map(preprocess))
  return datasets[split_map[split]].map(preprocess), train_mean
//  [python]
def compute_losses(self, obs):
    image = obs['image']
    action = obs['action']
    reward = obs['reward']
    state = self.initialize(tf.shape(image)[0])
    state['rewards'] = reward
    priors, posteriors = self.observe(action, image, state)
    features = self._get_features(posteriors)
    frames = self._predict_frames_tpl(features)
    if self._reward_from_frames:
      rewards = self._predict_reward_tpl(frames.mode(), reward[:, -1])
    else:
      rewards = self._predict_reward_tpl(features, reward[:, -1])
    obs_likelihood = frames.log_prob(image)
    reward_likelihood = rewards.log_prob(tf.to_float(reward))

    divergence = tfd.kl_divergence(
        self._get_distribution(posteriors), self._get_distribution(priors))
    divergence = tf.maximum(self._free_nats, divergence)
    loss = tf.reduce_mean(divergence - obs_likelihood -
                          reward_likelihood * self._reward_loss_mul)
    frames_mode = tf.clip_by_value((frames.mode() + 0.5) * 255, 0, 255)
    return (loss, tf.reduce_mean(reward_likelihood), tf.reduce_mean(divergence),
            tf.cast(frames_mode, dtype=tf.uint8), rewards.mode(), obs['reward'],
            tf.reduce_mean(tf.math.squared_difference(frames_mode, image)))
//  [python]
def compute_losses(self, obs):
    image = obs['image']
    action = obs['action']
    reward = obs['reward']
    state = self.initialize(tf.shape(image)[0])
    state['rewards'] = reward
    priors, posteriors = self.observe(action, image, state)
    features = self._get_features(posteriors)
    frames = self._predict_frames_tpl(features)
    if self._reward_from_frames:
      rewards = self._predict_reward_tpl(frames.mode(), reward[:, -1])
    else:
      rewards = self._predict_reward_tpl(features, reward[:, -1])
    obs_likelihood = frames.log_prob(image)
    reward_likelihood = rewards.log_prob(tf.cast(reward, dtype=tf.float32))

    divergence = tfd.kl_divergence(
        self._get_distribution(posteriors), self._get_distribution(priors))
    divergence = tf.maximum(self._free_nats, divergence)
    loss = tf.reduce_mean(divergence - obs_likelihood -
                          reward_likelihood * self._reward_loss_mul)
    frames_mode = tf.clip_by_value((frames.mode() + 0.5) * 255, 0, 255)
    return (loss, tf.reduce_mean(reward_likelihood), tf.reduce_mean(divergence),
            tf.cast(frames_mode, dtype=tf.uint8), rewards.mode(), obs['reward'],
            tf.reduce_mean(tf.math.squared_difference(frames_mode, image)))
//  [python]
def observe(self, actions, images, state):
    embedded_obs = self._encoder_tpl(images)
    if self.is_discrete_action:
      actions = tf.one_hot(
          actions[:, :, 0], self._action_space.n, dtype=tf.float32)
    else:
      actions = tf.to_float(actions)
    actions = tf.transpose(actions, [1, 0, 2])
    embedded_obs = tf.transpose(embedded_obs, [1, 0, 2])
    state.pop('rewards')
    priors, posteriors = static_scan(
        lambda prev, inp: self._posterior_tpl(prev[1], *inp),
        (actions, embedded_obs), (state, state))
    priors = {
        key: tf.transpose(value, [1, 0, 2]) for key, value in priors.items()
    }
    posteriors = {
        key: tf.transpose(value, [1, 0, 2])
        for key, value in posteriors.items()
    }
    return priors, posteriors
//  [python]
def observe(self, actions, images, state):
    embedded_obs = self._encoder_tpl(images)
    if self.is_discrete_action:
      actions = tf.one_hot(
          actions[:, :, 0], self._action_space.n, dtype=tf.float32)
    else:
      actions = tf.cast(actions, tf.float32)
    actions = tf.transpose(actions, [1, 0, 2])
    embedded_obs = tf.transpose(embedded_obs, [1, 0, 2])
    state.pop('rewards')
    priors, posteriors = static_scan(
        lambda prev, inp: self._posterior_tpl(prev[1], *inp),
        (actions, embedded_obs), (state, state))
    priors = {
        key: tf.transpose(value, [1, 0, 2]) for key, value in priors.items()
    }
    posteriors = {
        key: tf.transpose(value, [1, 0, 2])
        for key, value in posteriors.items()
    }
    return priors, posteriors
//  [python]
def _encoder(self, images):
    kwargs = dict(strides=2, activation=tf.nn.relu)
    images = tf.to_float(images)
    hidden = tf.reshape(images, [-1] + images.shape[2:].as_list())
    hidden = tf.layers.conv2d(hidden, 32, 4, **kwargs)
    hidden = tf.layers.conv2d(hidden, 64, 4, **kwargs)
    hidden = tf.layers.conv2d(hidden, 128, 4, **kwargs)
    hidden = tf.layers.conv2d(hidden, 256, 4, **kwargs)
    hidden = tf.layers.flatten(hidden)
    assert hidden.shape[1:].as_list() == [1024], hidden.shape.as_list()
    embedded_obs = tf.reshape(hidden, [
        tf.shape(images)[0],
        tf.shape(images)[1],
        np.prod(hidden.shape[1:].as_list())
    ])
    return embedded_obs
//  [python]
def _encoder(self, images):
    kwargs = dict(strides=2, activation=tf.nn.relu)
    images = tf.cast(images, tf.float32)
    hidden = tf.reshape(images, [-1] + images.shape[2:].as_list())
    hidden = tf.layers.conv2d(hidden, 32, 4, **kwargs)
    hidden = tf.layers.conv2d(hidden, 64, 4, **kwargs)
    hidden = tf.layers.conv2d(hidden, 128, 4, **kwargs)
    hidden = tf.layers.conv2d(hidden, 256, 4, **kwargs)
    hidden = tf.layers.flatten(hidden)
    assert hidden.shape[1:].as_list() == [1024], hidden.shape.as_list()
    embedded_obs = tf.reshape(hidden, [
        tf.shape(images)[0],
        tf.shape(images)[1],
        np.prod(hidden.shape[1:].as_list())
    ])
    return embedded_obs
//  [python]
def predict(self, actions, state):
    if isinstance(self._action_space, gym.spaces.Discrete):
      actions = tf.one_hot(
          actions[:, :, 0], self._action_space.n, dtype=tf.float32)
    else:
      actions = tf.to_float(actions)
    actions = tf.transpose(actions, [1, 0, 2])
    rewards = tf.to_float(state.pop('rewards'))
    priors = static_scan(self._prior_tpl, actions, state)
    priors = {
        key: tf.transpose(value, [1, 0, 2]) for key, value in priors.items()
    }
    features = self._get_features(priors)
    results = {}
    if self._reward_from_frames:
      frames = self._predict_frames_tpl(features).mode()
      results['reward'] = self._predict_reward_tpl(frames, rewards).mode()
    else:
      results['reward'] = self._predict_reward_tpl(features, rewards).mode()
    if self._include_frames_in_prediction:
      results['image'] = tf.cast(
          tf.clip_by_value(
              (self._predict_frames_tpl(features).mode() + 0.5) * 255, 0, 255),
          dtype=tf.uint8)
    return results
//  [python]
def predict(self, actions, state):
    if isinstance(self._action_space, gym.spaces.Discrete):
      actions = tf.one_hot(
          actions[:, :, 0], self._action_space.n, dtype=tf.float32)
    else:
      actions = tf.cast(actions, dtype=tf.float32)
    actions = tf.transpose(actions, [1, 0, 2])
    rewards = tf.cast(state.pop('rewards'), dtype=tf.float32)
    priors = static_scan(self._prior_tpl, actions, state)
    priors = {
        key: tf.transpose(value, [1, 0, 2]) for key, value in priors.items()
    }
    features = self._get_features(priors)
    results = {}
    if self._reward_from_frames:
      frames = self._predict_frames_tpl(features).mode()
      results['reward'] = self._predict_reward_tpl(frames, rewards).mode()
    else:
      results['reward'] = self._predict_reward_tpl(features, rewards).mode()
    if self._include_frames_in_prediction:
      results['image'] = tf.cast(
          tf.clip_by_value(
              (self._predict_frames_tpl(features).mode() + 0.5) * 255, 0, 255),
          dtype=tf.uint8)
    return results
//  [python]
def create_planet_observe_fn(model, model_dir, strategy):
  with strategy.scope():
    checkpoint = tf.train.Checkpoint(**model.get_trackables())
    manager = tf.train.CheckpointManager(checkpoint, model_dir, max_to_keep=1)
    checkpoint.restore(manager.latest_checkpoint).expect_partial()

  @tf.function
  def observe(images, actions, rewards, state):
    images = tf.to_float(images) / 255.0 - 0.5
    # break down the inputs along the batch dimension to form equal sized
    # tensors in each replica.
    num_replicas = strategy.num_replicas_in_sync
    images = tf.split(images, num_replicas)
    actions = tf.split(actions, num_replicas)
    state = {key: tf.split(value, num_replicas) for key, value in state.items()}
    devices = values.ReplicaDeviceMap(strategy.extended.worker_devices)
    dist_images = values.PerReplica(devices, tuple(images))
    dist_actions = values.PerReplica(devices, tuple(actions))
    dist_state = []
    for i in range(num_replicas):
      dist_state.append({key: value[i] for key, value in state.items()})
    dist_state = values.PerReplica(devices, tuple(dist_state))
    _, dist_posteriors = strategy.experimental_run_v2(
        model.observe, args=(dist_actions, dist_images, dist_state))
    dist_posteriors = {
        key: strategy.experimental_local_results(value)
        for key, value in dist_posteriors.items()
    }
    posteriors = {
        key: tf.concat(value, axis=0) for key, value in dist_posteriors.items()
    }
    posteriors = {key: value[:, -1] for key, value in posteriors.items()}
    posteriors['rewards'] = rewards[:, -1]
    return posteriors

  return observe
//  [python]
def create_planet_observe_fn(model, model_dir, strategy):
  with strategy.scope():
    checkpoint = tf.train.Checkpoint(**model.get_trackables())
    manager = tf.train.CheckpointManager(checkpoint, model_dir, max_to_keep=1)
    checkpoint.restore(manager.latest_checkpoint).expect_partial()

  @tf.function
  def observe(images, actions, rewards, state):
    images = tf.cast(images, tf.float32) / 255.0 - 0.5
    # break down the inputs along the batch dimension to form equal sized
    # tensors in each replica.
    num_replicas = strategy.num_replicas_in_sync
    images = tf.split(images, num_replicas)
    actions = tf.split(actions, num_replicas)
    state = {key: tf.split(value, num_replicas) for key, value in state.items()}
    devices = values.ReplicaDeviceMap(strategy.extended.worker_devices)
    dist_images = values.PerReplica(devices, tuple(images))
    dist_actions = values.PerReplica(devices, tuple(actions))
    dist_state = []
    for i in range(num_replicas):
      dist_state.append({key: value[i] for key, value in state.items()})
    dist_state = values.PerReplica(devices, tuple(dist_state))
    _, dist_posteriors = strategy.experimental_run_v2(
        model.observe, args=(dist_actions, dist_images, dist_state))
    dist_posteriors = {
        key: strategy.experimental_local_results(value)
        for key, value in dist_posteriors.items()
    }
    posteriors = {
        key: tf.concat(value, axis=0) for key, value in dist_posteriors.items()
    }
    posteriors = {key: value[:, -1] for key, value in posteriors.items()}
    posteriors['rewards'] = rewards[:, -1]
    return posteriors

  return observe
//  [python]
def testFactoredTensorImplicitConversion(self):
    a = np.random.rand(3, 4, 5)
    b = np.random.rand(6, 5)
    c = np.random.rand(3, 4, 6)
    # a factored representation of a Tensor of shape (3, 4, 6)
    factored = common_layers.FactoredTensor(tf.to_float(a), tf.to_float(b))
    # implicitly converts factored to a Tensor (performing the matmul)
    d = factored + tf.to_float(c)
    out = self.evaluate(d)
    self.assertEqual(out.shape, (3, 4, 6))
//  [python]
def testFactoredTensorImplicitConversion(self):
    a = np.random.rand(3, 4, 5)
    b = np.random.rand(6, 5)
    c = np.random.rand(3, 4, 6)
    # a factored representation of a Tensor of shape (3, 4, 6)
    factored = common_layers.FactoredTensor(tf.cast(a, dtype=tf.float32), tf.cast(b, dtype=tf.float32))
    # implicitly converts factored to a Tensor (performing the matmul)
    d = factored + tf.cast(c, dtype=tf.float32)
    out = self.evaluate(d)
    self.assertEqual(out.shape, (3, 4, 6))
//  [python]
def testPaddingCrossEntropyFactored(self):
    vocab_size = 19
    rows = 5
    cols = 4
    depth = 11
    label_smoothing = 0.1
    features = np.random.rand(rows, cols, depth)
    weights = np.random.rand(vocab_size, depth)
    labels = np.random.randint(0, vocab_size - 1, size=(rows, cols))
    with self.test_session() as session:
      features = tf.to_float(features)
      weights = tf.to_float(weights)
      labels = tf.to_int32(labels)
      logits = tf.matmul(
          tf.reshape(features, [rows * cols, depth]), weights, transpose_b=True)
      logits = tf.reshape(logits, [rows, cols, vocab_size])
      loss_num, loss_den = common_layers.padded_cross_entropy(
          logits, labels, label_smoothing=label_smoothing, reduce_sum=False)
      factored_logits = common_layers.FactoredTensor(features, weights)
      loss_num_f, loss_den_f = common_layers.padded_cross_entropy_factored(
          factored_logits,
          labels=labels,
          label_smoothing=label_smoothing,
          reduce_sum=False)
      num, den, num_f, den_f = session.run(
          [loss_num, loss_den, loss_num_f, loss_den_f])
    self.assertEqual(num.shape, (rows, cols))
    self.assertEqual(den.shape, (rows, cols))
    self.assertEqual(num_f.shape, (rows, cols))
    self.assertEqual(den_f.shape, (rows, cols))
    self.assertAllClose(num, num_f)
    self.assertAllClose(den, den_f)
//  [python]
def testPaddingCrossEntropyFactored(self):
    vocab_size = 19
    rows = 5
    cols = 4
    depth = 11
    label_smoothing = 0.1
    features = np.random.rand(rows, cols, depth)
    weights = np.random.rand(vocab_size, depth)
    labels = np.random.randint(0, vocab_size - 1, size=(rows, cols))
    with self.test_session() as session:
      features = tf.cast(features, tf.float32)
      weights = tf.cast(weights, tf.float32)
      labels = tf.cast(labels, tf.int32)
      logits = tf.matmul(
          tf.reshape(features, [rows * cols, depth]), weights, transpose_b=True)
      logits = tf.reshape(logits, [rows, cols, vocab_size])
      loss_num, loss_den = common_layers.padded_cross_entropy(
          logits, labels, label_smoothing=label_smoothing, reduce_sum=False)
      factored_logits = common_layers.FactoredTensor(features, weights)
      loss_num_f, loss_den_f = common_layers.padded_cross_entropy_factored(
          factored_logits,
          labels=labels,
          label_smoothing=label_smoothing,
          reduce_sum=False)
      num, den, num_f, den_f = session.run(
          [loss_num, loss_den, loss_num_f, loss_den_f])
    self.assertEqual(num.shape, (rows, cols))
    self.assertEqual(den.shape, (rows, cols))
    self.assertEqual(num_f.shape, (rows, cols))
    self.assertEqual(den_f.shape, (rows, cols))
    self.assertAllClose(num, num_f)
    self.assertAllClose(den, den_f)
//  [python]
def testPaddingCrossEntropyFactoredGrad(self):
    vocab_size = 19
    rows = 5
    cols = 4
    depth = 11
    label_smoothing = 0.1
    features = np.random.rand(rows, cols, depth)
    weights = np.random.rand(vocab_size, depth)
    labels = np.random.randint(0, vocab_size - 1, size=(rows, cols))
    with self.test_session() as session:
      features = tf.to_float(features)
      weights = tf.to_float(weights)
      labels = tf.to_int32(labels)
      logits = tf.matmul(
          tf.reshape(features, [rows * cols, depth]), weights, transpose_b=True)
      logits = tf.reshape(logits, [rows, cols, vocab_size])
      loss_num, loss_den = common_layers.padded_cross_entropy(
          logits, labels, label_smoothing=label_smoothing, reduce_sum=False)
      factored_logits = common_layers.FactoredTensor(features, weights)
      loss_num_factored, loss_den_factored = (
          common_layers.padded_cross_entropy_factored(
              factored_logits,
              labels=labels,
              label_smoothing=label_smoothing,
              reduce_sum=False))
      df, dw = tf.gradients(ys=[loss_num, loss_den], xs=[features, weights])
      df_factored, dw_factored = tf.gradients(
          ys=[loss_num_factored, loss_den_factored], xs=[features, weights])
      actual_df, actual_dw, actual_df_factored, actual_dw_factored = (
          session.run([df, dw, df_factored, dw_factored]))
    self.assertEqual(actual_df.shape, (rows, cols, depth))
    self.assertEqual(actual_dw.shape, (vocab_size, depth))
    self.assertEqual(actual_df_factored.shape, (rows, cols, depth))
    self.assertEqual(actual_dw_factored.shape, (vocab_size, depth))
    self.assertAllClose(actual_df, actual_df_factored)
    self.assertAllClose(actual_dw, actual_dw_factored)
//  [python]
def testPaddingCrossEntropyFactoredGrad(self):
    vocab_size = 19
    rows = 5
    cols = 4
    depth = 11
    label_smoothing = 0.1
    features = np.random.rand(rows, cols, depth)
    weights = np.random.rand(vocab_size, depth)
    labels = np.random.randint(0, vocab_size - 1, size=(rows, cols))
    with self.test_session() as session:
      features = tf.cast(features, tf.float32)
      weights = tf.cast(weights, tf.float32)
      labels = tf.cast(labels, tf.int32)
      logits = tf.matmul(
          tf.reshape(features, [rows * cols, depth]), weights, transpose_b=True)
      logits = tf.reshape(logits, [rows, cols, vocab_size])
      loss_num, loss_den = common_layers.padded_cross_entropy(
          logits, labels, label_smoothing=label_smoothing, reduce_sum=False)
      factored_logits = common_layers.FactoredTensor(features, weights)
      loss_num_factored, loss_den_factored = (
          common_layers.padded_cross_entropy_factored(
              factored_logits,
              labels=labels,
              label_smoothing=label_smoothing,
              reduce_sum=False))
      df, dw = tf.gradients(ys=[loss_num, loss_den], xs=[features, weights])
      df_factored, dw_factored = tf.gradients(
          ys=[loss_num_factored, loss_den_factored], xs=[features, weights])
      actual_df, actual_dw, actual_df_factored, actual_dw_factored = (
          session.run([df, dw, df_factored, dw_factored]))
    self.assertEqual(actual_df.shape, (rows, cols, depth))
    self.assertEqual(actual_dw.shape, (vocab_size, depth))
    self.assertEqual(actual_df_factored.shape, (rows, cols, depth))
    self.assertEqual(actual_dw_factored.shape, (vocab_size, depth))
    self.assertAllClose(actual_df, actual_df_factored)
    self.assertAllClose(actual_dw, actual_dw_factored)
//  [python]
def pairwise_dist_periodic(positions, periodicity):
    assert len(positions.shape) == 2
    with tf.name_scope("pairwise_dist_periodic"):
        vecs = periodic_distance(
            tf.expand_dims(positions, axis=1),
            tf.expand_dims(positions, axis=0),
            periodicity,
        )
        mask = tf.to_float(tf.equal(vecs, 0.0))
        vecs = vecs + mask * 1e-16  # gradient infinite for 0
        dists = tf.norm(vecs, axis=2)
        return dists
//  [python]
def pairwise_dist_periodic(positions, periodicity):
    assert len(positions.shape) == 2
    with tf.name_scope("pairwise_dist_periodic"):
        vecs = periodic_distance(
            tf.expand_dims(positions, axis=1),
            tf.expand_dims(positions, axis=0),
            periodicity,
        )
        mask = tf.cast(tf.equal(vecs, 0.0), tf.float32)
        vecs = vecs + mask * 1e-16  # gradient infinite for 0
        dists = tf.norm(vecs, axis=2)
        return dists
//  [python]
def testConvHiddenReluMemoryEfficient(self):
    batch = 3
    length = 23
    io_size = 16
    filter_size = 7
    x = np.random.rand(batch, length, io_size)
    dy = np.random.rand(batch, length, io_size)
    with self.test_session() as session:
      x = tf.to_float(x)
      dy = tf.to_float(dy)
      f1 = tf.get_variable("f1", [1, io_size, filter_size])
      f2 = tf.get_variable("f2", [1, filter_size, io_size])
      norm_scale, norm_bias = common_layers.layer_norm_vars(io_size)
      y = common_layers.conv_hidden_relu_memory_efficient(
          x, filter_size, forget=False,
          test_vars=(f1, f2, norm_scale, norm_bias))
      y_forget = common_layers.conv_hidden_relu_memory_efficient(
          x, filter_size, forget=True,
          test_vars=(f1, f2, norm_scale, norm_bias))
      dx, df1, df2, dnorm_scale, dnorm_bias = tf.gradients(
          ys=[y], xs=[x, f1, f2, norm_scale, norm_bias], grad_ys=[dy])
      dx_f, df1_f, df2_f, dnorm_scale_f, dnorm_bias_f = tf.gradients(
          ys=[y_forget], xs=[x, f1, f2, norm_scale, norm_bias], grad_ys=[dy])
      session.run(tf.global_variables_initializer())
      (y, y_forget,
       dx, df1, df2, dnorm_scale, dnorm_bias,
       dx_f, df1_f, df2_f, dnorm_scale_f, dnorm_bias_f) = session.run(
           [y, y_forget,
            dx, df1, df2, dnorm_scale, dnorm_bias,
            dx_f, df1_f, df2_f, dnorm_scale_f, dnorm_bias_f])
    self.assertAllClose(y, y_forget)
    self.assertAllClose(df2, df2_f, rtol=2e-6, atol=2e-6)
    self.assertAllClose(df1, df1_f, rtol=2e-6, atol=2e-6)
    self.assertAllClose(dnorm_scale, dnorm_scale_f)
    self.assertAllClose(dnorm_bias, dnorm_bias_f)
    self.assertAllClose(dx, dx_f)
//  [python]
def testConvHiddenReluMemoryEfficient(self):
    batch = 3
    length = 23
    io_size = 16
    filter_size = 7
    x = np.random.rand(batch, length, io_size)
    dy = np.random.rand(batch, length, io_size)
    with self.test_session() as session:
      x = tf.cast(x, tf.float32)
      dy = tf.cast(dy, tf.float32)
      f1 = tf.get_variable("f1", [1, io_size, filter_size])
      f2 = tf.get_variable("f2", [1, filter_size, io_size])
      norm_scale, norm_bias = common_layers.layer_norm_vars(io_size)
      y = common_layers.conv_hidden_relu_memory_efficient(
          x, filter_size, forget=False,
          test_vars=(f1, f2, norm_scale, norm_bias))
      y_forget = common_layers.conv_hidden_relu_memory_efficient(
          x, filter_size, forget=True,
          test_vars=(f1, f2, norm_scale, norm_bias))
      dx, df1, df2, dnorm_scale, dnorm_bias = tf.gradients(
          ys=[y], xs=[x, f1, f2, norm_scale, norm_bias], grad_ys=[dy])
      dx_f, df1_f, df2_f, dnorm_scale_f, dnorm_bias_f = tf.gradients(
          ys=[y_forget], xs=[x, f1, f2, norm_scale, norm_bias], grad_ys=[dy])
      session.run(tf.global_variables_initializer())
      (y, y_forget,
       dx, df1, df2, dnorm_scale, dnorm_bias,
       dx_f, df1_f, df2_f, dnorm_scale_f, dnorm_bias_f) = session.run(
           [y, y_forget,
            dx, df1, df2, dnorm_scale, dnorm_bias,
            dx_f, df1_f, df2_f, dnorm_scale_f, dnorm_bias_f])
    self.assertAllClose(y, y_forget)
    self.assertAllClose(df2, df2_f, rtol=2e-6, atol=2e-6)
    self.assertAllClose(df1, df1_f, rtol=2e-6, atol=2e-6)
    self.assertAllClose(dnorm_scale, dnorm_scale_f)
    self.assertAllClose(dnorm_bias, dnorm_bias_f)
    self.assertAllClose(dx, dx_f)
//  [python]
def get_NER_indices(self, top_span_starts, top_span_ends, NER_starts, NER_ends, NER_indices):
    N_starts = tf.concat([top_span_starts, NER_starts], 0)
    N_ends = tf.concat([top_span_ends, NER_ends], 0)
    max_idx = tf.reduce_max([N_starts, N_ends])
    s_oh = tf.one_hot(top_span_starts, max_idx)
    label_s_oh = tf.one_hot(NER_starts, max_idx)
    e_oh = tf.one_hot(top_span_ends, max_idx)
    label_e_oh = tf.one_hot(NER_ends, max_idx)

    s_mult = tf.matmul(s_oh, label_s_oh, transpose_b = True)
    e_mult = tf.matmul(e_oh, label_e_oh, transpose_b = True)

    output = tf.reduce_max(s_mult*e_mult, axis = 1)

    output2 = tf.matmul(s_mult*e_mult, tf.cast(tf.expand_dims(NER_indices, -1), tf.float32))
    output2 = tf.squeeze(output2, 1)
    output2 = tf.to_int32(output2)
    output2 = tf.one_hot(output2, 6)

    return output2, output
//  [python]
def get_NER_indices(self, top_span_starts, top_span_ends, NER_starts, NER_ends, NER_indices):
    N_starts = tf.concat([top_span_starts, NER_starts], 0)
    N_ends = tf.concat([top_span_ends, NER_ends], 0)
    max_idx = tf.reduce_max([N_starts, N_ends])
    s_oh = tf.one_hot(top_span_starts, max_idx)
    label_s_oh = tf.one_hot(NER_starts, max_idx)
    e_oh = tf.one_hot(top_span_ends, max_idx)
    label_e_oh = tf.one_hot(NER_ends, max_idx)

    s_mult = tf.matmul(s_oh, label_s_oh, transpose_b = True)
    e_mult = tf.matmul(e_oh, label_e_oh, transpose_b = True)

    output = tf.reduce_max(s_mult*e_mult, axis = 1)

    output2 = tf.matmul(s_mult*e_mult, tf.cast(tf.expand_dims(NER_indices, -1), tf.float32))
    output2 = tf.squeeze(output2, 1)
    output2 = tf.cast(output2, tf.int32)
    output2 = tf.one_hot(output2, 6)

    return output2, output
//  [python]
def get_mention_indices(self, top_span_starts, top_span_ends, entity_starts, entity_ends, entity_indices):
    N_starts = tf.concat([top_span_starts, entity_starts], 0)
    N_ends = tf.concat([top_span_ends, entity_ends], 0)
    max_idx = tf.reduce_max([N_starts, N_ends])
    s_oh = tf.one_hot(top_span_starts, max_idx)
    label_s_oh = tf.one_hot(entity_starts, max_idx)
    e_oh = tf.one_hot(top_span_ends, max_idx)
    label_e_oh = tf.one_hot(entity_ends, max_idx)

    s_mult = tf.matmul(s_oh, label_s_oh, transpose_b = True)
    e_mult = tf.matmul(e_oh, label_e_oh, transpose_b = True)

    output = tf.reduce_max(s_mult*e_mult, axis=1)
    output2 = tf.matmul(s_mult*e_mult, tf.cast(tf.expand_dims(entity_indices, -1), tf.float32))
    output2 = tf.squeeze(output2, 1)
    output2 = tf.to_int32(output2)

    return output2
//  [python]
def get_mention_indices(self, top_span_starts, top_span_ends, entity_starts, entity_ends, entity_indices):
    N_starts = tf.concat([top_span_starts, entity_starts], 0)
    N_ends = tf.concat([top_span_ends, entity_ends], 0)
    max_idx = tf.reduce_max([N_starts, N_ends])
    s_oh = tf.one_hot(top_span_starts, max_idx)
    label_s_oh = tf.one_hot(entity_starts, max_idx)
    e_oh = tf.one_hot(top_span_ends, max_idx)
    label_e_oh = tf.one_hot(entity_ends, max_idx)

    s_mult = tf.matmul(s_oh, label_s_oh, transpose_b = True)
    e_mult = tf.matmul(e_oh, label_e_oh, transpose_b = True)

    output = tf.reduce_max(s_mult*e_mult, axis=1)
    output2 = tf.matmul(s_mult*e_mult, tf.cast(tf.expand_dims(entity_indices, -1), tf.float32))
    output2 = tf.squeeze(output2, 1)
    output2 = tf.cast(output2, tf.int32)

    return output2
//  [python]
def cross_entropy_seq_with_mask(logits, target_seqs, input_mask, return_details=False, name=None):
    """Returns the expression of cross-entropy of two sequences, implement
    softmax internally. Normally be used for Dynamic RNN outputs.

    Parameters
    -----------
    logits : network identity outputs
        2D tensor, ``network.outputs``, [batch_size, number of output units].
    target_seqs : int of tensor, like word ID.
        [batch_size, ?]
    input_mask : the mask to compute loss
        The same size with target_seqs, normally 0 and 1.
    return_details : boolean
        - If False (default), only returns the loss.
        - If True, returns the loss, losses, weights and targets (reshape to one vetcor).

    Examples
    --------
    - see Image Captioning Example.
    """
    targets = tf.reshape(target_seqs, [-1])   # to one vector
    weights = tf.to_float(tf.reshape(input_mask, [-1]))   # to one vector like targets
    losses = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=targets, name=name) * weights
    #losses = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=targets, name=name)) # for TF1.0 and others

    try: ## TF1.0
        loss = tf.divide(tf.reduce_sum(losses),   # loss from mask. reduce_sum before element-wise mul with mask !!
                        tf.reduce_sum(weights),
                        name="seq_loss_with_mask")
    except: ## TF0.12
        loss = tf.div(tf.reduce_sum(losses),   # loss from mask. reduce_sum before element-wise mul with mask !!
                        tf.reduce_sum(weights),
                        name="seq_loss_with_mask")
    if return_details:
        return loss, losses, weights, targets
    else:
        return loss