code: # coding=utf-8
  # Copyright 2022 The Google Research Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Forward processing of raw data to sRGB images.
  
  Unprocessing Images for Learned Raw Denoising
  http://timothybrooks.com/tech/unprocessing
  """
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import tensorflow.compat.v1 as tf
  
  
  def apply_gains(bayer_images, red_gains, blue_gains):
    """Applies white balance gains to a batch of Bayer images."""
    bayer_images.shape.assert_is_compatible_with((None, None, None, 4))
    green_gains = tf.ones_like(red_gains)
    gains = tf.stack([red_gains, green_gains, green_gains, blue_gains], axis=-1)
    gains = gains[:, tf.newaxis, tf.newaxis, :]
    return bayer_images * gains
  
  
  def demosaic(bayer_images):
    """Bilinearly demosaics a batch of RGGB Bayer images."""
    bayer_images.shape.assert_is_compatible_with((None, None, None, 4))
  
    # This implementation exploits how edges are aligned when upsampling with
    # tf.image.resize_bilinear().
  
    with tf.name_scope(None, 'demosaic'):
      shape = tf.shape(bayer_images)
      shape = [shape[1] * 2, shape[2] * 2]
  
      red = bayer_images[Ellipsis, 0:1]
      red = tf.image.resize_bilinear(red, shape)
  
      green_red = bayer_images[Ellipsis, 1:2]
      green_red = tf.image.flip_left_right(green_red)
      green_red = tf.image.resize_bilinear(green_red, shape)
      green_red = tf.image.flip_left_right(green_red)
      green_red = tf.space_to_depth(green_red, 2)
  
      green_blue = bayer_images[Ellipsis, 2:3]
      green_blue = tf.image.flip_up_down(green_blue)
      green_blue = tf.image.resize_bilinear(green_blue, shape)
      green_blue = tf.image.flip_up_down(green_blue)
      green_blue = tf.space_to_depth(green_blue, 2)
  
      green_at_red = (green_red[Ellipsis, 0] + green_blue[Ellipsis, 0]) / 2
      green_at_green_red = green_red[Ellipsis, 1]
      green_at_green_blue = green_blue[Ellipsis, 2]
      green_at_blue = (green_red[Ellipsis, 3] + green_blue[Ellipsis, 3]) / 2
  
      green_planes = [
          green_at_red, green_at_green_red, green_at_green_blue, green_at_blue
      ]
      green = tf.depth_to_space(tf.stack(green_planes, axis=-1), 2)
  
      blue = bayer_images[Ellipsis, 3:4]
      blue = tf.image.flip_up_down(tf.image.flip_left_right(blue))
      blue = tf.image.resize_bilinear(blue, shape)
      blue = tf.image.flip_up_down(tf.image.flip_left_right(blue))
  
      rgb_images = tf.concat([red, green, blue], axis=-1)
      return rgb_images
  
  
  def apply_ccms(images, ccms):
    """Applies color correction matrices."""
    images.shape.assert_has_rank(4)
    images = images[:, :, :, tf.newaxis, :]
    ccms = ccms[:, tf.newaxis, tf.newaxis, :, :]
    return tf.reduce_sum(images * ccms, axis=-1)
  
  
  def gamma_compression(images, gamma=2.2):
    """Converts from linear to gamma space."""
    # Clamps to prevent numerical instability of gradients near zero.
    return tf.maximum(images, 1e-8) ** (1.0 / gamma)
  
  
  def process(bayer_images, red_gains, blue_gains, cam2rgbs):
    """Processes a batch of Bayer RGGB images into sRGB images."""
    bayer_images.shape.assert_is_compatible_with((None, None, None, 4))
    with tf.name_scope(None, 'process'):
      # White balance.
      bayer_images = apply_gains(bayer_images, red_gains, blue_gains)
      # Demosaic.
      bayer_images = tf.clip_by_value(bayer_images, 0.0, 1.0)
      images = demosaic(bayer_images)
      # Color correction.
      images = apply_ccms(images, cam2rgbs)
      # Gamma compression.
      images = tf.clip_by_value(images, 0.0, 1.0)
      images = gamma_compression(images)
    return images
   # coding=utf-8
  # Copyright 2022 The Google Research Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Tests for ...object_detection.preprocessor."""
  
  import tensorflow as tf
  from tf3d import standard_fields
  from tf3d.object_detection import preprocessor
  
  
  class ObjectDetectionPreprocessorTest(tf.test.TestCase):
  
    def _image_correspondence_fn(self, inputs):
      return {
          'view_images': {
              'rgb_view':
                  tf.cast(
                      tf.zeros([5, 200, 300, 3], dtype=tf.int32), dtype=tf.uint8),
          },
          'view_indices_2d': {
              'rgb_view':
                  tf.random.uniform([5, 100, 2],
                                    minval=-10,
                                    maxval=1000,
                                    dtype=tf.int32)
          }
      }
  
    def _get_input_dict(self, height=240, width=320):
      return {
          standard_fields.InputDataFields.camera_image:
              tf.zeros((height, width, 3), dtype=tf.uint8),
          standard_fields.InputDataFields.point_positions:
              tf.random.uniform((100, 3), minval=-1, maxval=1),
          standard_fields.InputDataFields.camera_intrinsics:
              tf.constant([
                  [160.0, 0.0, 160.0],  # fx,  s, cx
                  [0.0, 160.0, 120.0],  #  0, fy, cy
                  [0.0, 0.0, 1.0],  #  0,  0,  1
              ]),
          standard_fields.InputDataFields.camera_rotation_matrix:
              tf.eye(3),
          standard_fields.InputDataFields.camera_translation:
              tf.constant([0., 0., 2.]),
          standard_fields.InputDataFields.objects_class:
              tf.constant([1, 4, 5]),
          standard_fields.InputDataFields.objects_length:
              tf.constant([[4.0], [1.0], [1.0]]),
          standard_fields.InputDataFields.objects_height:
              tf.constant([[2.0], [1.0], [4.0]]),
          standard_fields.InputDataFields.objects_width:
              tf.constant([[2.0], [1.0], [1.0]]),
          standard_fields.InputDataFields.objects_rotation_matrix:
              tf.constant([[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
                           [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
                           [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]]),
          standard_fields.InputDataFields.objects_center:
              tf.constant([[4.0, 4.0, 4.0], [2.5, 2.5, 2.5], [0.5, 1.5, 9.5]]),
          standard_fields.InputDataFields.objects_difficulty:
              tf.constant([[1], [1], [1]]),
          standard_fields.InputDataFields.objects_instance_id:
              tf.constant([[1], [2], [1]]),
          standard_fields.InputDataFields.objects_has_3d_info:
              tf.constant([1, 1, 0]),
          standard_fields.InputDataFields.camera_image_name:
              tf.convert_to_tensor('image', tf.string),
      }
  
    def test_preprocess_output_shapes(self):
      height, width = (240, 320)
      input_dict = self._get_input_dict(height, width)
      object_keys = preprocessor._OBJECT_KEYS
      output_keys = [
          standard_fields.InputDataFields.camera_intrinsics,
          standard_fields.InputDataFields.camera_rotation_matrix,
          standard_fields.InputDataFields.camera_translation,
          standard_fields.InputDataFields.point_positions,
          standard_fields.InputDataFields.num_valid_points,
          standard_fields.InputDataFields.object_class_points,
          standard_fields.InputDataFields.object_center_points,
          standard_fields.InputDataFields.object_height_points,
          standard_fields.InputDataFields.object_width_points,
          standard_fields.InputDataFields.object_rotation_matrix_points,
          standard_fields.InputDataFields.object_length_points,
          standard_fields.InputDataFields.object_instance_id_points,
      ]
      output_dict = preprocessor.preprocess(
          inputs=input_dict,
          images_points_correspondence_fn=self._image_correspondence_fn,
          image_preprocess_fn_dic=None)
      for key in output_keys:
        self.assertIn(key, output_dict)
      self.assertEqual(
          output_dict[standard_fields.InputDataFields.camera_intrinsics].shape,
          (3, 3))
      self.assertEqual(
          output_dict[
              standard_fields.InputDataFields.camera_rotation_matrix].shape,
          (3, 3))
      self.assertEqual(
          output_dict[standard_fields.InputDataFields.camera_translation].shape,
          (3,))
      self.assertEqual(
          output_dict[standard_fields.InputDataFields.point_positions].shape,
          (100, 3))
      self.assertEqual(
          output_dict[standard_fields.InputDataFields.num_valid_points].numpy(),
          100)
      self.assertEqual(
          output_dict[standard_fields.InputDataFields.object_class_points].shape,
          (100,))
      self.assertEqual(
          output_dict[standard_fields.InputDataFields.object_center_points].shape,
          (100, 3))
      self.assertEqual(
          output_dict[standard_fields.InputDataFields.object_height_points].shape,
          (100, 1))
      self.assertEqual(
          output_dict[standard_fields.InputDataFields.object_width_points].shape,
          (100, 1))
      self.assertEqual(
          output_dict[standard_fields.InputDataFields.object_length_points].shape,
          (100, 1))
      self.assertEqual(
          output_dict[standard_fields.InputDataFields
                      .object_rotation_matrix_points].shape, (100, 3, 3))
      self.assertEqual(
          output_dict[
              standard_fields.InputDataFields.object_instance_id_points].shape,
          (100,))
      for key in object_keys:
        self.assertEqual(output_dict[key].shape[0], 2)
  
    def test_preprocess_output_keys(self):
      height, width = (240, 320)
      input_dict = self._get_input_dict(height, width)
      output_dict = preprocessor.preprocess(
          inputs=input_dict,
          images_points_correspondence_fn=self._image_correspondence_fn,
          output_keys=[standard_fields.InputDataFields.camera_image],
          image_preprocess_fn_dic=None)
      self.assertIn(standard_fields.InputDataFields.camera_image, output_dict)
      self.assertEqual(len(output_dict.keys()), 1)
  
    def test_preprocess_missing_input_raises(self):
      with self.assertRaises(ValueError):
        empty_input = {}
        preprocessor.preprocess(inputs=empty_input)
  
  
  if __name__ == '__main__':
    tf.test.main()
   # coding=utf-8
  # Copyright 2022 The Google Research Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  # pylint: skip-file
  import tensorflow.compat.v1 as tf
  import math
  from combiner.tf import attention
  from combiner.tf import ops
  import functools
  
  
  def shift_right(x, axis):
    """Shift input x to the right along given axis."""
    pad_widths = [(0, 0)] * len(x.shape)
    pad_widths[axis] = (1, 0)
    padded = tf.pad(x, pad_widths)
    return tf.slice(padded, begin=[0]*len(x.shape), size=x.shape)
  
  
  def shift_left(x, axis):
    """Shift input x to the left along given axis."""
    pad_widths = [(0, 0)] * len(x.shape)
    pad_widths[axis] = (0, 1)
    padded = tf.pad(x, pad_widths)
    begin = [0]*len(x.shape)
    begin[axis] = 1
    return tf.slice(padded, begin=begin, size=x.shape)
  
  
  def approx_cummax(x, axis, exclusive=False, reverse=False):
    """Approximate the cummax operation in jax."""
    sum_x = tf.math.cumsum(x, axis, exclusive=exclusive, reverse=reverse)
    # return tf.math.cumsum(tf.nn.relu(x), axis, reverse=reverse)
    return sum_x
  
  
  def get_causal_mask(x, axis, is_strict, upper=False):
    """Get attention mask bias (keep a lower triangle).
  
    Args:
      x: input tensor
      axis: across which dim to make mask
      is_strict: if True, the diagonal will be masked out as well.
      upper: upper or lower triangle
  
    Returns:
      mask: tensor of {0, -1e9} ^ (x.shape[axis], x.shape[axis])
    """
    seq_len = tf.shape(x)[axis]
    if is_strict:
      if upper:
        mask = tf.linalg.band_part(
            tf.ones([seq_len, seq_len], dtype=x.dtype),
            num_lower=-1, num_upper=0)
      else:
        mask = tf.linalg.band_part(
            tf.ones([seq_len, seq_len], dtype=x.dtype),
            num_lower=0, num_upper=-1)
    else:
      if upper:
        mask = 1.0 - tf.linalg.band_part(
            tf.ones([seq_len, seq_len], dtype=x.dtype),
            num_lower=0, num_upper=-1)
      else:
        mask = 1.0 - tf.linalg.band_part(
            tf.ones([seq_len, seq_len], dtype=x.dtype),
            num_lower=-1, num_upper=0)
    mask = -1e9 * mask
    return mask
  
  
  def pooling_summary(x, axis, local_summary, keepdims=False):
    """Perform a cheap pooling summary of a span.
  
    Args:
      x: input tensor
      axis: over which axis to summarize
      local_summary: str of format activation-pooling, choose
        from {relu, identity}-{max, sum, mean}
      keepdims: whether to keep the summarized singleton axis.
  
    Returns:
      y: the same shape as x for other axis,
        except y.shape[axis] = 1 if keepdims=True,
        otherwise y.rank = x.rank + 1
    """
    act, pool = local_summary.split('-')
    if act == 'relu':
      x = tf.nn.relu(x)
    elif act == 'identity':
      pass
    elif act == 'deepset':
      x = ops.trail_dense(x, x.shape.as_list()[-1], bias=False)
      x = tf.nn.relu(x)
    else:
      raise ValueError('Unsupported activation: %s' % act)
    if pool == 'mean':
      x = tf.math.reduce_mean(x, axis=axis, keepdims=keepdims)
    elif pool == 'max':
      x = tf.math.reduce_max(x, axis=axis, keepdims=keepdims)
    elif pool == 'sum':
      x = tf.math.reduce_sum(x, axis=axis, keepdims=keepdims)
    else:
      raise ValueError('Unsupported pooling: %s' % pool)
    return x
  
  
  def axial_mixture_unidir(x, config, is_training=True, causal=True):
    """Full attention matrix with axial pattern as local and mixture for global summary."""
    del is_training
    assert causal
    bsize = x.shape[0]
    query, key, value = attention.get_qkv(x, x, x, hidden_size=config.model_size,
                                          num_heads=config.num_heads, bias=config.dense_use_bias)
    head_dim = config.model_size // config.num_heads
    assert config.max_seq_len % config.max_seg_len == 0
    num_seg = config.max_seq_len // config.max_seg_len
    cur_query = tf.reshape(query, [bsize,
                                   num_seg,
                                   config.max_seg_len,
                                   config.num_heads,
                                   head_dim])
    cur_key = tf.reshape(key, cur_query.shape)
    cur_val = tf.reshape(value, cur_query.shape)
  
    col_logit_expr = 'BSUNK,BTUNK->BUNST'
    col_attn_expr = 'BUNST,BTUNK->BSUNK'
    col_strict_mask = get_causal_mask(cur_query,
                               axis=1,
                               is_strict=True)[tf.newaxis, tf.newaxis, tf.newaxis, :, :]
    row_logit_expr = 'BUSNK,BUTNK->BUNST'
    row_attn_expr = 'BUNST,BUTNK->BUSNK'
    row_mask = get_causal_mask(cur_query,
                               axis=2,
                               is_strict=False)[tf.newaxis, tf.newaxis, tf.newaxis, :, :]
    col_logits = tf.einsum(col_logit_expr, cur_query, cur_key) + col_strict_mask
    row_logits = tf.einsum(row_logit_expr, cur_query, cur_key) + row_mask
  
    ###################
    col_up2down_query = approx_cummax(cur_query, axis=1)
    col_up2down_key = shift_right(approx_cummax(cur_key, axis=1), axis=1)
    col_mask = get_causal_mask(
        cur_query, axis=1, is_strict=False)[tf.newaxis, tf.newaxis,
                                            tf.newaxis, :, :]
    col_up2down_logits = tf.einsum(col_logit_expr, col_up2down_query,
                                   cur_key) + col_mask
    col_up2down_attn_weights = attention.float32_softmax(
        col_up2down_logits, axis=-1)
    col_up2down_summary = tf.einsum(col_attn_expr, col_up2down_attn_weights,
                                    cur_val)
    col_up2down_summary = shift_right(col_up2down_summary, axis=1)
  
    row_only_myself_mask = tf.eye(tf.shape(cur_query)[2], dtype=cur_query.dtype)[tf.newaxis, tf.newaxis, tf.newaxis, :, :]
    row_without_myself_mask = -1e9 * row_only_myself_mask
    all_maskout = tf.cast(tf.fill(row_without_myself_mask.shape, -1e9), cur_query.dtype)
    row_without_myself_mask = tf.concat([all_maskout] + [row_without_myself_mask] * (cur_query.shape[1] - 1),
                                        axis=1)
    previous_row_logits = tf.einsum(row_logit_expr, cur_query, col_up2down_key) + row_without_myself_mask
    ###################
  
    row_left2right_query = approx_cummax(cur_query, axis=2)
    row_left2right_key = shift_right(approx_cummax(cur_key, axis=2), axis=2)
    row_left2right_logits = tf.einsum(row_logit_expr, row_left2right_query,
                                      cur_key) + row_mask
    row_left2right_attn_weights = attention.float32_softmax(
        row_left2right_logits, axis=-1)
    row_left2right_summary = tf.einsum(row_attn_expr, row_left2right_attn_weights,
                                       cur_val)
    row_left2right_summary = shift_right(row_left2right_summary, axis=2)
  
    all_maskout = tf.cast(tf.fill(col_strict_mask.shape, -1e9), cur_query.dtype)
    col_strict_without_first_mask = tf.concat(
        [all_maskout] + [col_strict_mask] * (cur_query.shape[2] - 1), axis=1)
    top_left_col_logits = tf.einsum(
        col_logit_expr, cur_query,
        row_left2right_key) + col_strict_without_first_mask
    ###################
  
    row_right2left_query = approx_cummax(cur_query, axis=2, reverse=True)
    row_right2left_key = shift_left(
        approx_cummax(cur_key, axis=2, reverse=True), axis=2)
    row_upper_mask = get_causal_mask(
        cur_query, axis=2, is_strict=False, upper=True)[tf.newaxis, tf.newaxis,
                                                        tf.newaxis, :, :]
    row_right2left_logits = tf.einsum(row_logit_expr, row_right2left_query,
                                      cur_key) + row_upper_mask
    row_right2left_attn_weights = attention.float32_softmax(
        row_right2left_logits, axis=-1)
    row_right2left_summary = tf.einsum(row_attn_expr, row_right2left_attn_weights,
                                       cur_val)
    row_right2left_summary = shift_left(row_right2left_summary, axis=2)
    col_strict_without_last_mask = tf.concat(
        [col_strict_mask] * (cur_query.shape[2] - 1) + [all_maskout], axis=1)
    top_right_col_logits = tf.einsum(
        col_logit_expr, cur_query,
        row_right2left_key) + col_strict_without_last_mask
    ###################
  
    joint_logits = tf.concat([
        tf.transpose(col_logits, perm=[0, 3, 2, 1, 4]), row_logits,
        previous_row_logits,
        tf.transpose(top_left_col_logits, perm=[0, 3, 2, 1, 4]),
        tf.transpose(top_right_col_logits, perm=[0, 3, 2, 1, 4])
    ],
                             axis=-1)
    attn_weights = attention.float32_softmax(joint_logits, axis=-1)
    col_att, row_att, previous_row_att, top_left_col_att, top_right_col_att = tf.split(attn_weights,
                                                                                       [num_seg,
                                                                                        config.max_seg_len,
                                                                                        config.max_seg_len,
                                                                                        num_seg,
                                                                                        num_seg], axis=-1)
    col_att = tf.transpose(col_att, [0, 3, 2, 1, 4])
    top_left_col_att = tf.transpose(top_left_col_att, [0, 3, 2, 1, 4])
    top_right_col_att = tf.transpose(top_right_col_att, [0, 3, 2, 1, 4])
    col_merged = tf.einsum(col_attn_expr, col_att, cur_val)
    row_merged = tf.einsum(row_attn_expr, row_att, cur_val)
    previous_row_merged = tf.einsum(row_attn_expr, previous_row_att,
                                    col_up2down_summary)
    top_left_merged = tf.einsum(col_attn_expr, top_left_col_att,
                                row_left2right_summary)
    top_right_merged = tf.einsum(col_attn_expr, top_right_col_att,
                                 row_right2left_summary)
  
    joint_merged = tf.reshape(
        col_merged + row_merged + previous_row_merged + top_left_merged +
        top_right_merged,
        [bsize, num_seg * config.max_seg_len, config.num_heads, head_dim])
    output = ops.trail_dense(joint_merged, config.model_size, begin_axis=-2)
    return output
  
  
  def sqrt_fixed_full(x, config, is_training=True, causal=True):
    """Full attention matrix with sqrt decomposition."""
    bsize = x.shape[0]
    query, key, value = attention.get_qkv(x, x, x, hidden_size=config.model_size,
                                          num_heads=config.num_heads,
                                          bias=config.dense_use_bias)
    head_dim = config.model_size // config.num_heads
    assert config.max_seq_len % config.max_seg_len == 0
    num_seg = config.max_seq_len // config.max_seg_len
    cur_query = tf.reshape(query, [-1,
                                   num_seg,
                                   config.max_seg_len,
                                   config.num_heads,
                                   head_dim])
    with tf.variable_scope('pooling_query'):
      merged_query = pooling_summary(cur_query, axis=2,
                                     local_summary=config.local_summary,
                                     keepdims=True)
    cur_key = tf.reshape(key, cur_query.shape)
    cur_val = tf.reshape(value, cur_query.shape)
    span_val = attention.dot_product_attention(merged_query,
                                               cur_key,
                                               cur_val,
                                               is_training=is_training,
                                               attn_axis=1,
                                               dropatt=config.dropatt)
    span_val = tf.squeeze(span_val, axis=2)
    with tf.variable_scope('pooling_key'):
      span_key = pooling_summary(cur_key, axis=2,
                                 local_summary=config.local_summary,
                                 keepdims=False)
    local_logits = tf.einsum('bsqhd,bskhd->bsqhk', cur_query, cur_key)
    if causal:
      local_mask = get_causal_mask(cur_query, axis=2, is_strict=False)
      local_mask = tf.expand_dims(local_mask, axis=-2)
      local_logits += local_mask
    prev_logits = tf.einsum('bqhd,bkhd->bqhk', query, span_key)
    if causal:
      prev_mask = get_causal_mask(cur_query, axis=1, is_strict=True)
      prev_mask = tf.repeat(prev_mask, [config.max_seg_len] * num_seg, axis=0)
      prev_logits += tf.expand_dims(prev_mask, axis=1)
    joint_logits = tf.concat([tf.reshape(local_logits,
                                         [bsize, config.max_seq_len,
                                          config.num_heads, -1]),
                              prev_logits], axis=-1)
    attn_weights = attention.float32_softmax(joint_logits, axis=-1)
    local_att, prev_att = tf.split(attn_weights, [config.max_seg_len, num_seg],
                                   axis=-1)
    if is_training:
      local_att = tf.nn.dropout(local_att, rate=config.dropatt)
    local_att = tf.reshape(local_att, [bsize, num_seg,
                                       config.max_seg_len,
                                       config.num_heads,
                                       config.max_seg_len])
    local_merged = tf.einsum('bsqhk,bskhd->bsqhd', local_att, cur_val)
    prev_merged = tf.einsum('bqhk,bkhd->bqhd', prev_att, span_val)
    joint_merged = prev_merged + tf.reshape(local_merged, prev_merged.shape)
    output = ops.trail_dense(joint_merged, config.model_size, begin_axis=-2)
    return output
  
  
  def axial_rowmajor(x, config, is_training=True, causal=True):
    """Full attention matrix with sqrt decomposition."""
    bsize = x.shape[0]
    seq_len = x.shape.as_list()[1]
    head_dim = config.model_size // config.num_heads
    assert seq_len % config.max_seg_len == 0
    num_seg = seq_len // config.max_seg_len
    x_sqr = tf.reshape(x,
                       [bsize, num_seg, config.max_seg_len, config.model_size])
    q_row_local, key_row_local, value_row_local = attention.get_qkv(
        x_sqr, x_sqr, x_sqr, hidden_size=config.model_size,
        num_heads=config.num_heads, bias=config.dense_use_bias)
    local_logits = tf.einsum('bsqhd,bskhd->bsqhk', q_row_local, key_row_local)
    row_probs = attention.float32_softmax(local_logits, axis=-1)
    if is_training:
      row_probs = tf.nn.dropout(row_probs, rate=config.dropatt)
  
    row_attn_out = tf.einsum('bsqhk,bskhd->bsqhd', row_probs, value_row_local)
    if config.row_summary == 'none':
      key_row = key_row_local
    elif config.row_summary in ['wsum', 'proj', 'wsum_proj']:
      if 'wsum' in config.row_summary:
        pre_summary = tf.einsum('bsqhk,bskhd->bsqhd', row_probs, key_row_local)
      else:
        pre_summary = row_attn_out
      if 'proj' in config.row_summary:
        with tf.variable_scope('rowmajor_param_post'):
          key_row = ops.trail_dense(pre_summary, config.model_size, begin_axis=-2,
                                    bias=config.dense_use_bias)
          key_row = ops.postprocess(x_sqr, key_row, config, is_training)
          _, key_row = ops.preprocess(key_row, config)
          key_row = ops.trail_dense(key_row, [config.num_heads, head_dim],
                                    bias=config.dense_use_bias)
      else:
        key_row = pre_summary
    else:
      raise ValueError('Unknown row summary %s' % config.row_summary)
    if causal:
      local_mask = get_causal_mask(q_row_local, axis=2, is_strict=False)
      local_logits += local_mask[:, tf.newaxis, :]
  
    global_logits = tf.einsum('bqlhd,bklhd->bqlhk', q_row_local, key_row)
    if causal:
      global_mask = get_causal_mask(q_row_local, axis=1, is_strict=True)
      global_logits += global_mask[:, tf.newaxis, tf.newaxis, :]
    # (bsize, num_seg, seg_len, n_head, seg_len + num_seg)
    joint_logits = tf.concat([local_logits, global_logits], axis=-1)
    attn_probs = attention.float32_softmax(joint_logits, axis=-1)
    local_att, global_att = tf.split(attn_probs,
                                     [config.max_seg_len, num_seg],
                                     axis=-1)
    if is_training:
      local_att = tf.nn.dropout(local_att, rate=config.dropatt)
    local_merged = tf.einsum('bsqhk,bskhd->bsqhd', local_att, value_row_local)
    global_merged = tf.einsum('bqlhv,bvlhd->bqlhd', global_att, row_attn_out)
    joint_merged = tf.reshape(local_merged + global_merged,
                              [bsize, seq_len,
                               config.num_heads, head_dim])
    output = ops.trail_dense(joint_merged, config.model_size,
                             begin_axis=-2, bias=config.dense_use_bias)
    return output
  
  
  def axial_mixture_bidir(x, config, is_training=True, causal=False):
    """Full attention matrix with axial mixture decomposition."""
    assert not causal
    bsize = x.shape[0]
    seq_len = x.shape.as_list()[1]
    head_dim = config.model_size // config.num_heads
    assert seq_len % config.max_seg_len == 0
    num_seg = seq_len // config.max_seg_len
    x_sqr = tf.reshape(x,
                       [bsize, num_seg, config.max_seg_len, config.model_size])
    query, key, value = attention.get_qkv(
        x_sqr, x_sqr, x_sqr, hidden_size=config.model_size,
        num_heads=config.num_heads, bias=config.dense_use_bias)
    local_row_logits = tf.einsum('bushd,buthd->bhust', query, key)
    local_col_logits = tf.einsum('bsuhd,btuhd->bhsut', query, key)
    # TODO: add self-mask for local_col_logits
  
    span_attn_fn = functools.partial(attention.dot_product_attention,
                                     key_heads=key,
                                     value_heads=value,
                                     is_training=is_training,
                                     dropatt=config.dropatt)
  
    # === top-down summary ===
    col_query_topdown = approx_cummax(query, 1, exclusive=True)
    col_key_topdown = approx_cummax(key, 1, exclusive=True)
    col_t2d_mask = get_causal_mask(x_sqr, axis=1, is_strict=True)
    col_t2d_val = span_attn_fn(query_heads=col_query_topdown,
                               attn_axis=0,
                               attn_bias=col_t2d_mask)
  
    # === bottom-up summary ===
    col_query_bottomup = approx_cummax(query, 1, exclusive=True, reverse=True)
    col_key_bottomup = approx_cummax(key, 1, exclusive=True, reverse=True)
    col_b2t_mask = get_causal_mask(x_sqr, axis=1, is_strict=True, upper=True)
    col_b2t_val = span_attn_fn(query_heads=col_query_bottomup,
                               attn_axis=0,
                               attn_bias=col_b2t_mask)
  
    # === left2right summary ===
    row_query_left2right = approx_cummax(query, 2, exclusive=True)
    row_key_left2right = approx_cummax(key, 2, exclusive=True)
    row_l2r_mask = get_causal_mask(x_sqr, axis=2, is_strict=True)
    row_l2r_val = span_attn_fn(query_heads=row_query_left2right,
                               attn_axis=1,
                               attn_bias=row_l2r_mask)
  
    # === right2left summary ===
    row_query_right2left = approx_cummax(query, 2, exclusive=True, reverse=True)
    row_key_right2left = approx_cummax(key, 2, exclusive=True, reverse=True)
    row_r2l_mask = get_causal_mask(x_sqr, axis=2, is_strict=True, upper=True)
    row_r2l_val = span_attn_fn(query_heads=row_query_right2left,
                               attn_axis=1,
                               attn_bias=row_r2l_mask)
  
    global_t2d_logits = tf.einsum('bushd,buthd->bhust', query, col_key_topdown)
    global_b2t_logits = tf.einsum('bushd,buthd->bhust', query, col_key_bottomup)
    global_l2r_logits = tf.einsum('bsuhd,btuhd->bhsut', query, row_key_left2right)
    global_r2l_logits = tf.einsum('bsuhd,btuhd->bhsut', query, row_key_right2left)
    joint_logits = tf.concat([local_row_logits, local_col_logits,
                              global_t2d_logits, global_b2t_logits,
                              global_l2r_logits, global_r2l_logits], axis=-1)
    attn_probs = attention.float32_softmax(joint_logits, axis=-1)
    prow, pcol, pt2d, pb2t, pl2r, pr2l = tf.split(
        attn_probs, [config.max_seg_len, num_seg, config.max_seg_len,
                     config.max_seg_len, num_seg, num_seg], axis=-1)
    mrow = tf.einsum('bhust,buthd->bushd', prow, value)
    mcol = tf.einsum('bhsut,btuhd->bsuhd', pcol, value)
    mt2d = tf.einsum('bhust,buthd->bushd', pt2d, col_t2d_val)
    mb2t = tf.einsum('bhust,buthd->bushd', pb2t, col_b2t_val)
    ml2r = tf.einsum('bhsut,btuhd->bsuhd', pl2r, row_l2r_val)
    mr2l = tf.einsum('bhsut,btuhd->bsuhd', pr2l, row_r2l_val)
    joint_merged = mrow + mcol + mt2d + mb2t + ml2r + mr2l
    joint_merged = tf.reshape(joint_merged,
                              [bsize, seq_len, config.num_heads, head_dim])
    output = ops.trail_dense(joint_merged, config.model_size,
                             begin_axis=-2, bias=config.dense_use_bias)
    return output
   # coding=utf-8
  # Copyright 2022 The Google Research Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Input function hook for PPO TF estimator.
  
  For the PPO algorithm, see https://arxiv.org/abs/1707.06347.
  """
  from absl import logging
  import gin
  import numpy as np
  import tensorflow.compat.v1 as tf
  from polish.ppo import ppo_loss
  from polish.utils import distributions
  from polish.utils import host_call_fn
  from polish.utils import tf_layers
  from tensorflow.contrib import tpu as contrib_tpu
  
  logging.set_verbosity(logging.INFO)
  
  
  @gin.configurable
  class PpoModelFn(object):
    """Main class for model function used in tf.estimator.
  
    Attributes:
      policy_loss: Proximal Policy Optimization (PPO) policy loss.
      value_loss: PPO value loss.
      entropy_loss: PPO entropy loss.
      imitation_kl_divergence: The KL-divergence of action distributions between
        the policy and MCTS.
      total_loss: PPO total loss.
      clipfrac: Fraction of examples in a batch clipped by PPO.
      approxkl: `Approximate` KL divergence between new policy and old policy.
        This is an estimate (approximate) of the KL divergence, since we compute
        the KL divergence using the samples drawn from the new and
        old distributions.
      total_params: Total trainable parameters.
      train_op: Training operation.
      mean_new: Mean of new policy distribution.
      logstd_new: Log standard deviation of new policy distribution.
      mean_old: Mean of old policy distribution.
      logstd_old: Log standard deviation of old policy distribution.
      value_new: state-value from the latest trained state-value network.
      kl_divergence: Kullback-Leibler divergence between new and old policy.
      entropy: Entropy of the new policy.
      global_step: Global step of training.
      policy_ratio: the ratio between new policy and old policy.
      last_iteration_mcts_enable: Track the sampling type (PPO sampling/MCTS) in
        the process of training. That is, whether in the last iteration of
        training, we use PPO sampling (False) or MCTS sampling (True).
      mcts_sampling_enable: If True, it means that the current batch
          is generated by MCTS.
      mean_mse_loss: Mean squared error between the mean value of
        policy distribuiton and the mean value returned by MCTS given a state.
      logstd_mse_loss: Mean squared error between log of standard deviation value
        of the policy distribuiton and the log of standard deviation value
        returned by MCTS given a state.
    """
  
    def __init__(
        self,
        env_action_space=2,
        iterations_per_loop=320,
        num_timesteps=1000000,
        max_horizon=2048,
        learning_rate=3e-4,
        use_tpu=False,
        ppo2_enable=True,
        policy_coeff=1.0,
        value_coeff=0.5,
        entropy_coeff=0.0,
        tpu_num_shards=8,
        mse_loss_coeff=0.0,
        warmstart_file=None,
        policy_hidden_layer_size=64,
        value_hidden_layer_size=64):
      """Creates a model function for PPO algorithm.
  
      The default values for all the parameters are from PPO paper.
  
      Args:
        env_action_space: The size of environment action space.
        iterations_per_loop: Number of steps to run on TPU before outfeeding
          metrics to the CPU. If the number of iterations in the loop would exceed
          the number of train steps, the loop will exit before reaching
          --iterations_per_loop. The larger this value is, the higher the
          utilization on the TPU.
        num_timesteps: Total number of timesteps. Defines the total number of
          samples taken from the environment during the whole process of training.
        max_horizon: Maximum number of samples taken from the environment before
          starting training.
        learning_rate: Initial learning rate value. Note that the actual learning
          rate is linearly decayed.
        use_tpu: If True, training occurs on TPU.
        ppo2_enable: If True, we use the next version of PPO algorithm, known as
          PPO2. In this version, not only does the probability ratio get clipped,
          but also the clipping is performed on the value loss.
          For more information:
          https://github.com/openai/baselines/tree/master/baselines/ppo2.
        policy_coeff: Policy loss coefficient in the total loss calculation.
        value_coeff: Value loss coefficient in the total loss calculation.
        entropy_coeff: Entropy loss coefficient in the total loss calculation.
        tpu_num_shards: Number of TPU shards.
        mse_loss_coeff: The coefficient for Mean Squared Error (MSE) loss.
        warmstart_file: If not None, we restore the weights for the parameters in
          `newpolicy` scope from this file. `newpolicy` scope contains both
          policy and value network.
        policy_hidden_layer_size: The size of hidden layer in policy network.
          Currently, this value is used for both of the hidden layers.
        value_hidden_layer_size: The size of hidden layer in value network.
          Currently, this value is used for both of the hidden layers.
      """
      self.policy_loss = 0
      self.value_loss = 0
      self.entropy_loss = 0
      self.total_loss = 0
      self.clipfrac = 0
      self.approxkl = 0
      self.policy_ratio = 0
  
      self.total_params = None
      self.train_op = None
  
      self.mean_new = None
      self.logstd_new = None
      self.mean_old = None
      self.logstd_old = None
      self.value_new = None
  
      self.kl_divergence = None
      self.entropy = 0
  
      self.global_step = None
  
      self._decayed_learning_rate = None
  
      self._env_action_space = env_action_space
      self._iterations_per_loop = iterations_per_loop
      self._num_timesteps = num_timesteps
      self._max_horizon = max_horizon
      self._learning_rate = learning_rate
      self._use_tpu = use_tpu
      self._ppo2_enable = ppo2_enable
      self._policy_coeff = policy_coeff
      self._value_coeff = value_coeff
      self._entropy_coeff = entropy_coeff
      self._tpu_num_shards = tpu_num_shards
  
      self._mse_loss_coeff = mse_loss_coeff
      self._warmstart_file = warmstart_file
  
      self.last_iteration_mcts_enable = False
      self._mcts_global_step = 0
  
      self._policy_hidden_layer_size = policy_hidden_layer_size
      self._value_hidden_layer_size = value_hidden_layer_size
  
    def __call__(self, features, labels, mode, params):
      return self.model_fn(features, labels, mode, params)
  
    def model_inference_fn_ppo(self, features, prefix):
      """Builds just the inference part of the model graph.
  
      Args:
        features: input features tensor.
        prefix: prefix to be added to the network.
  
      Returns:
        (value, var, mean) tuple of tensors.
      """
      # Policy Network
      features = tf.layers.flatten(features)
      with tf.variable_scope(prefix + 'policy', reuse=tf.AUTO_REUSE):
        policy_1 = tf.tanh(
            tf_layers.fc(
                tensor_in=features,
                num_hidden=self._policy_hidden_layer_size,
                scope_name='/policy_1',
                init_scale=np.sqrt(2)))
        policy_2 = tf.tanh(
            tf_layers.fc(
                tensor_in=policy_1,
                num_hidden=self._policy_hidden_layer_size,
                scope_name='/policy_2',
                init_scale=np.sqrt(2)))
        mean = tf_layers.fc(
            tensor_in=policy_2,
            num_hidden=self._env_action_space,
            scope_name='/mean',
            init_scale=0.01,
            init_bias=0.0)
        logstd_var = tf.get_variable(
            name=prefix + '_logstd',
            shape=[1, self._env_action_space],
            initializer=tf.zeros_initializer())
        # Evaluate logstd_var and broadcast to have a same shape as mean
        logstd = tf.multiply(logstd_var, 1.0)
  
        value_1 = tf.tanh(
            tf_layers.fc(
                tensor_in=features,
                num_hidden=self._value_hidden_layer_size,
                scope_name='/value_1',
                init_scale=np.sqrt(2)))
        value_2 = tf.tanh(
            tf_layers.fc(
                tensor_in=value_1,
                num_hidden=self._value_hidden_layer_size,
                scope_name='/value_2',
                init_scale=np.sqrt(2)))
        value = tf_layers.fc(
            tensor_in=value_2, num_hidden=1, scope_name='/value')[:, 0]
  
      return value, logstd, mean
  
    def learning_rate_update_true_fn(self):
      """The function which is performed if the predicate is true.
  
      The predicate that calls this function is defined in 'update_learning_rate'.
  
      Returns:
        The current global step.
      """
      return tf.train.get_global_step()
  
    def learning_rate_update_false_fn(self):
      """The function which is performed if the predicate is false.
  
      The predicate that calls this function is defined in 'update_learning_rate'.
  
      Returns:
        A type-casted value of `_mcts_global_step` to int64.
        `_mcts_global_step` is the global step at which MCTS algorithm starts.
        The type casting is necessary as the type of returned tensor in `true_fn`
        is an int.64.
      """
      return tf.cast(self._mcts_global_step, tf.int64)
  
    def update_learning_rate(self):
      """Update the learning rate with a decaying factor.
      """
      self._current_global_step = tf.cond(
          tf.equal(self.mcts_sampling_enable,
                   True), lambda: self._mcts_global_step, lambda: 0)
  
      self._current_global_step = tf.cast(self._current_global_step, tf.int64)
  
      update = (tf.train.get_global_step() -
                self._current_global_step) // self._iterations_per_loop + 1
      current_frac = self._num_timesteps // self._max_horizon
      update = tf.cast(update, tf.float32)
      current_frac = tf.cast(current_frac, tf.float32)
      frac = 1.0 - (update - 1.0) / current_frac
      self._decayed_learning_rate = self._learning_rate * frac
      self._mcts_global_step = tf.cond(
          tf.not_equal(self.mcts_sampling_enable,
                       self.last_iteration_mcts_enable),
          self.learning_rate_update_true_fn, self.learning_rate_update_false_fn)
      self.last_iteration_mcts_enable = self.mcts_sampling_enable
  
    def build_training_op(self, loss):
      """Get training operation.
  
      Args:
        loss: a loss function for training.
  
      Define the optimization operation and perform gradient calculation for both
        TPU/Non-TPU training.
  
      Returns:
        Computed gradient.
      """
      adam_optimizer = tf.train.AdamOptimizer(
          learning_rate=self._decayed_learning_rate, epsilon=1e-5)
      if self._use_tpu:
        # If we use TPUs, reduce_mean runs on each chip separately and by default
        # only the loss of the first chip is reported.
        #
        # You can either:
        # - execute this if, which synchronizes the losses
        #   across the chips to obtain the full loss on all samples.
        # - or remove this section, gaining some performance and getting the
        #   loss only from the first chip.
        # compute gradients perform averaging of the loss
        adam_optimizer = tf.tpu.CrossShardOptimizer(adam_optimizer)
  
        tpu_sum_loss = contrib_tpu.cross_replica_sum(loss / self._tpu_num_shards)
  
        grads_and_vars = adam_optimizer.compute_gradients(tpu_sum_loss,
                                                          self.total_params)
        grads, var = zip(*grads_and_vars)
        sum_grads = []
        sum_vars = []
        for (grad, var) in grads_and_vars:
          if grad is None:
            sum_grads.append(grad)
            sum_vars.append(var)
          else:
            sum_grads.append(
                contrib_tpu.cross_replica_sum(grad) / self._tpu_num_shards)
            sum_vars.append(var)
        # calculate sum of grads
        norm_grads, _ = tf.clip_by_global_norm(sum_grads, 0.5)
        grads_and_vars = list(zip(norm_grads, sum_vars))
      else:
        grads_and_vars = adam_optimizer.compute_gradients(loss,
                                                          self.total_params)
        grads, var = zip(*grads_and_vars)
        norm_grads, _ = tf.clip_by_global_norm(grads, 0.5)
        grads_and_vars = list(zip(norm_grads, var))
  
      return adam_optimizer.apply_gradients(
          grads_and_vars, global_step=tf.train.get_global_step())
  
    def calc_normalized_advantage(self, return_tensor, value_tensor):
      """Compute General Advantage Estimation (GAE) and normalize it.
  
      Note that, the advantage calculation-normalization is performed for a batch
        of data.
  
      Args:
        return_tensor: The discounted accumulated reward (return) calculated
          for the given rollout trajectory.
        value_tensor: The value for each state for the given rollout trajectory.
  
      Returns:
        Returns the normalized General Advantage Estimation (GAE).
      """
      batch_advantage = return_tensor - value_tensor
      batch_advantage_std = tf.keras.backend.std(batch_advantage)
      batch_advantage_mean = tf.reduce_mean(batch_advantage)
      batch_advantage_norm = (batch_advantage - batch_advantage_mean) / (
          batch_advantage_std + 1e-8)
      return batch_advantage_norm
  
    def create_host_call_fn(self, params):
      """Create host call function.
  
      `host_call` function is later called by TPU estimator to
        send some metrics to host for logging.
  
      Args:
        params: A dictionary of hyperparameters passed to the tf.estimator.
  
      Returns:
        A host call function that generates a set of tf summaries.
      """
      names_and_tensors = [
          ('Batch_Params/mean_mse_loss', self.mean_mse_loss),
          ('Batch_Params/logstd_mse_loss', self.logstd_mse_loss),
          ('Batch_Params/policy_loss', self.policy_loss),
          ('Batch_Params/mcts_enable', self.mcts_sampling_enable),
          ('Batch_Params/value_loss', self.value_loss),
          ('Batch_Params/policy_entropy', self.entropy_loss),
          ('Batch_Params/imitation_kl_divergence', self.imitation_kl_divergence),
          ('Batch_Params/clip_fraction', self.clipfrac),
          ('Batch_Params/max_ratio', tf.reduce_max(self.policy_ratio)),
          ('Batch_Params/min_ratio', tf.reduce_min(self.policy_ratio)),
          ('Batch_Params/mean_ratio', tf.reduce_mean(self.policy_ratio)),
          ('Batch_Params/approx_kl', self.approxkl),
          ('Learning_Rate/learning_rate', self._decayed_learning_rate),
          ('Learning_Rate/global_step', tf.train.get_global_step())
      ]
  
      return host_call_fn.build_host_call_fn_every_n_global_steps(
          params=params,
          names_and_tensors=names_and_tensors,
          n=self._iterations_per_loop)
  
    def compute_total_loss(self, pd_new, pd_old, value_tensor, return_tensor,
                           batch_advantage_norm,
                           policy_old_neg_logprob_tensor,
                           policy_action_tensor):
      """Defines the total loss function.
  
      Args:
        pd_new: The current policy distribution
          (a multivariate normal distribution). This policy distribution gets
          updated in the course of training.
        pd_old: The old policy distribution that we use during sampling the
          trajectory (a multivariate normal distribution).
        value_tensor: The values associated to the rollout trajectory.
        return_tensor: The return values computed for the rollout trajectory.
        batch_advantage_norm: The normalized advantage tensor computed for a
          batch of data. For advantage calculation, we use generalized
          advantage estimation (GAE) formula.
        policy_old_neg_logprob_tensor: The negative log probabilities from the
          policy rollouts.
        policy_action_tensor: The actions from the policy rollouts.
      """
      # Policy loss
      ppo_policy_loss_out = ppo_loss.ppo_policy_loss(
          neg_logprobs_old=policy_old_neg_logprob_tensor,
          actions=policy_action_tensor,
          advantages=batch_advantage_norm,
          dist_new=pd_new,
          mcts_sampling=self.mcts_sampling_enable)
  
      (self.policy_loss, self.approxkl, self.clipfrac,
       self.policy_ratio) = ppo_policy_loss_out
  
      # Value Loss
      if self._ppo2_enable:
        self.value_loss = ppo_loss.ppo2_value_loss(
            value_old=value_tensor,
            pred_value=self.value_new,
            returns=return_tensor)
      else:
        self.value_loss = ppo_loss.ppo1_value_loss(
            pred_value=self.value_new, returns=return_tensor)
  
      # MSE loss between mean and standard deviations
      self.mean_mse_loss, self.logstd_mse_loss = ppo_loss.l2_norm_policy_loss(
          policy_mean=self.mean_new,
          policy_logstd=self.logstd_new,
          mcts_mean=self.mean_old,
          mcts_logstd=self.logstd_old)
  
      mcts_dist = distributions.MultiVariateNormalDiag(
          mean=self.mean_old, logstd=self.logstd_old)
      policy_dist = distributions.MultiVariateNormalDiag(
          mean=self.mean_new, logstd=self.logstd_new)
      self.imitation_kl_divergence = tf.reduce_mean(
          policy_dist.kl_divergence(mcts_dist))
      # Calculate KL divergence and entropy of new distribution
      self.kl_divergence = tf.reduce_mean(pd_new.kl_divergence(pd_old))
      self.entropy = pd_new.entropy()
  
      # Calculate entropy loss
      self.entropy_loss = tf.reduce_mean(self.entropy)
  
      # Calulate total loss
      total_loss_ppo = (self._policy_coeff * self.policy_loss) + (
          self._value_coeff * self.value_loss) - (
              self._entropy_coeff * self.entropy_loss)
  
      total_loss_mcts = (self._value_coeff * self.value_loss) + (
          self._mse_loss_coeff *
          (self.imitation_kl_divergence + self.entropy_loss))
  
      self.total_loss = tf.cond(
          tf.equal(self.mcts_sampling_enable,
                   True), lambda: total_loss_mcts, lambda: total_loss_ppo)
  
    def model_fn(self, features, labels, mode, params):
      """The implementation of PPO algorithm.
  
      Args:
        features: dict from string to tensor with shape
              'state_tensor': [BATCH_SIZE, env.state_space]
        labels: dict from string to tensor with shape
                'action_tensor': [BATCH_SIZE, self._env_action_space]
                'advantage_tensor': [BATCH_SIZE]
                'returns_tensor': [BATCH_SIZE]
        mode: a tf.estimator.ModeKeys (batchnorm params update for TRAIN only).
        params: (Ignored; needed for compat with TPUEstimator).
  
      Returns:
        tf.estimator.EstimatorSpec with props.
        mode: same as mode arg.
        predictions: dict of tensors
                'mean': [BATCH_SIZE, self._env_action_space]
                'logstd': [BATCH_SIZE, self._env_action_space]
                'value': [BATCH_SIZE]
                'action': [BATCH_SIZE, self._env_action_space]
                'neg_logprob': [BATCH_SIZE, self._env_action_space]
        loss: a single value tensor.
        train_op: train op eval_metric_ops return dict of tensors.
      """
  
      # Policy network
      network_out = self.model_inference_fn_ppo(features['mcts_features'], 'new')
      self.value_new = network_out[0]
      self.logstd_new = network_out[1]
      self.mean_new = network_out[2]
  
      self.global_step = tf.train.get_or_create_global_step()
      # Sample an action
      pd_new = distributions.MultiVariateNormalDiag(
          mean=self.mean_new, logstd=self.logstd_new)
      action_sample = pd_new.sample()
      action_sample_neg_logprob = pd_new.negative_log_prob(action_sample)
  
      # Used during TF estimator prediction
      if mode == tf.estimator.ModeKeys.PREDICT:
        predictions = {
            'mean': self.mean_new,
            'logstd': self.logstd_new,
            'value': self.value_new,
            'action': action_sample,
            'neg_logprob': action_sample_neg_logprob
        }
        pred_estimator = tf.estimator.tpu.TPUEstimatorSpec(
            mode,
            predictions=predictions,
            export_outputs={
                'ppo_inference':
                    tf.estimator.export.PredictOutput({
                        'mean': self.mean_new,
                        'logstd': self.logstd_new,
                        'value': self.value_new,
                        'action': action_sample,
                        'neg_logprob': action_sample_neg_logprob
                    })
            })
        return pred_estimator.as_estimator_spec()
  
      # Placeholder
      self.mcts_sampling_enable = tf.reduce_all(labels['mcts_enable_tensor'])
  
      self.mean_old = labels['mean_tensor']
      self.logstd_old = labels['logstd_tensor']
      pd_old = distributions.MultiVariateNormalDiag(
          mean=self.mean_old, logstd=self.logstd_old)
  
      batch_advantage_norm = self.calc_normalized_advantage(
          return_tensor=labels['policy_return_tensor'],
          value_tensor=labels['policy_value_tensor'])
  
      self.compute_total_loss(pd_new, pd_old, labels['value_tensor'],
                              labels['return_tensor'], batch_advantage_norm,
                              labels['policy_old_neg_logprob_tensor'],
                              labels['policy_action_tensor'])
      # Update learning rate
      self.update_learning_rate()
  
      # Build training operation
      self.total_params = tf.trainable_variables(scope='newpolicy')
  
      train_ops = self.build_training_op(self.total_loss)
  
      host_call = self.create_host_call_fn(params)
  
      if mode != tf.estimator.ModeKeys.TRAIN:
        raise ValueError('Estimator mode should be train at this point.')
  
      if mode == tf.estimator.ModeKeys.TRAIN:
        # Setup fine tune scaffold
        # The scaffold here is used to restore the weights from _warmstart_file.
        # If _warmstart_file is None, the training starts from the beginning.
        if self._warmstart_file:
          logging.info('Warmstart')
          def tpu_scaffold():
            # restore all the variables
            tf.init_from_checkpoint(self._warmstart_file,
                                    {'newpolicy/': 'newpolicy/'})
            return tf.train.Scaffold()
          scaffold_fn = tpu_scaffold
        else:
          scaffold_fn = None
  
        tpu_estimator_spec = tf.estimator.tpu.TPUEstimatorSpec(
            mode=mode,
            loss=self.total_loss,
            train_op=train_ops,
            host_call=host_call,
            scaffold_fn=scaffold_fn)
        if self._use_tpu:
          return tpu_estimator_spec
        else:
          return tpu_estimator_spec.as_estimator_spec()
   # coding=utf-8
  # Copyright 2022 The Google Research Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  # Copyright 2022 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  # Usage: configure.py [--quiet] [--no-deps]
  #
  # Options:
  #  --quiet  Give less output.
  #  --no-deps  Don't install Python dependencies
  """Configures ScaNN to be built from source."""
  
  from __future__ import print_function
  
  import argparse
  import logging
  import os
  import subprocess
  import sys
  
  _BAZELRC = ".bazelrc"
  _BAZEL_QUERY = ".bazel-query.sh"
  
  
  # Writes variables to bazelrc file
  def write_to_bazelrc(line):
    with open(_BAZELRC, "a") as f:
      f.write(line + "\n")
  
  
  def write_action_env(var_name, var):
    write_to_bazelrc('build --action_env %s="%s"' % (var_name, str(var)))
    with open(_BAZEL_QUERY, "a") as f:
      f.write('{}="{}" '.format(var_name, var))
  
  
  def get_input(question):
    try:
      return input(question)
    except EOFError:
      return ""
  
  
  def generate_shared_lib_name(namespec):
    """Converts the linkflag namespec to the full shared library name."""
    # Assume Linux for now
    return namespec[1][3:]
  
  
  def create_build_configuration():
    """Main function to create build configuration."""
    print()
    print("Configuring ScaNN to be built from source...")
  
    pip_install_options = ["--upgrade"]
    parser = argparse.ArgumentParser()
    parser.add_argument("--quiet", action="store_true", help="Give less output.")
    parser.add_argument(
        "--no-deps",
        action="store_true",
        help="Do not check and install Python dependencies.",
    )
    args = parser.parse_args()
    if args.quiet:
      pip_install_options.append("--quiet")
  
    python_path = sys.executable
    with open("requirements.txt") as f:
      required_packages = f.read().splitlines()
  
    print()
    if args.no_deps:
      print("> Using pre-installed Tensorflow.")
    else:
      print("> Installing", required_packages)
      install_cmd = [python_path, "-m", "pip", "install"]
      install_cmd.extend(pip_install_options)
      install_cmd.extend(required_packages)
      subprocess.check_call(install_cmd)
  
    if os.path.isfile(_BAZELRC):
      os.remove(_BAZELRC)
    if os.path.isfile(_BAZEL_QUERY):
      os.remove(_BAZEL_QUERY)
  
    logging.disable(logging.WARNING)
  
    import tensorflow.compat.v2 as tf  # pylint: disable=g-import-not-at-top
  
    # pylint: disable=invalid-name
    _TF_CFLAGS = tf.sysconfig.get_compile_flags()
    _TF_LFLAGS = tf.sysconfig.get_link_flags()
    _TF_CXX11_ABI_FLAG = tf.sysconfig.CXX11_ABI_FLAG
  
    _TF_SHARED_LIBRARY_NAME = generate_shared_lib_name(_TF_LFLAGS)
    _TF_HEADER_DIR = _TF_CFLAGS[0][2:]
    _TF_SHARED_LIBRARY_DIR = _TF_LFLAGS[0][2:]
    # pylint: enable=invalid-name
  
    write_action_env("TF_HEADER_DIR", _TF_HEADER_DIR)
    write_action_env("TF_SHARED_LIBRARY_DIR", _TF_SHARED_LIBRARY_DIR)
    write_action_env("TF_SHARED_LIBRARY_NAME", _TF_SHARED_LIBRARY_NAME)
    write_action_env("TF_CXX11_ABI_FLAG", _TF_CXX11_ABI_FLAG)
  
    write_to_bazelrc("build --spawn_strategy=standalone")
    write_to_bazelrc("build --strategy=Genrule=standalone")
    write_to_bazelrc("build -c opt")
  
    print()
    print("Build configurations successfully written to", _BAZELRC)
    print()
  
    with open(_BAZEL_QUERY, "a") as f:
      f.write('bazel query "$@"')
  
  
  if __name__ == "__main__":
    create_build_configuration()
   # coding=utf-8
  # Copyright 2022 The Google Research Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  from __future__ import print_function
  from __future__ import division
  from __future__ import absolute_import
  
  import matplotlib
  matplotlib.use('Agg')
  import matplotlib.pyplot as plt
  import numpy as np
  
  
  def plot_heatmap(pdf_func, out_name, size=3):
    w = 100
    x = np.linspace(-size, size, w)
    y = np.linspace(-size, size, w)
    xx, yy = np.meshgrid(x, y)
    coords = np.stack([xx.flatten(), yy.flatten()]).transpose()
  
    scores = pdf_func(coords)
    a = scores.reshape((w, w))
  
    plt.imshow(a)
    plt.axis('equal')
    plt.axis('off')
    plt.savefig(out_name, bbox_inches='tight')
    plt.close()
  
  
  def plot_samples(samples, out_name):
    plt.scatter(samples[:, 0], samples[:, 1])
    plt.axis('equal')
    plt.savefig(out_name, bbox_inches='tight')
    plt.close()
  
  
  def plot_joint(dataset, samples, out_name):
    x = np.max(dataset)
    y = np.max(-dataset)
    z = np.ceil(max((x, y)))
    plt.scatter(dataset[:, 0], dataset[:, 1], c='r', marker='x')
    plt.scatter(samples[:, 0], samples[:, 1], c='b', marker='.')
    plt.legend(['training data', 'ADE sampled'])
    plt.axis('equal')
    plt.xlim(-z, z)
    plt.ylim(-z, z)
    plt.savefig(out_name, bbox_inches='tight')
    plt.close()
  
    fname = out_name.split('/')[-1]
    out_name = '/'.join(out_name.split('/')[:-1]) + '/none-' + fname
    plt.figure(figsize=(8, 8))
    plt.scatter(dataset[:, 0], dataset[:, 1], c='r', marker='x')
    plt.scatter(samples[:, 0], samples[:, 1], c='b', marker='.')
    plt.axis('equal')
    plt.xlim(-z, z)
    plt.ylim(-z, z)
    plt.savefig(out_name, bbox_inches='tight')
    plt.close()
  
apis: tensorflow.compat.v1.image.flip_left_right tensorflow.compat.v1.ones_like tensorflow.compat.v1.image.flip_up_down tensorflow.compat.v1.space_to_depth tensorflow.compat.v1.image.resize_bilinear tensorflow.compat.v1.concat tensorflow.compat.v1.maximum tensorflow.compat.v1.reduce_sum tensorflow.compat.v1.shape tensorflow.compat.v1.stack tensorflow.compat.v1.clip_by_value tensorflow.compat.v1.name_scope
apis: tensorflow.zeros tensorflow.eye tensorflow.convert_to_tensor tensorflow.random.uniform tensorflow.constant tensorflow.test.main
apis: tensorflow.compat.v1.math.reduce_mean tensorflow.compat.v1.expand_dims tensorflow.compat.v1.shape tensorflow.compat.v1.repeat tensorflow.compat.v1.nn.relu tensorflow.compat.v1.pad tensorflow.compat.v1.reshape tensorflow.compat.v1.transpose tensorflow.compat.v1.slice tensorflow.compat.v1.math.reduce_sum tensorflow.compat.v1.math.reduce_max tensorflow.compat.v1.ones tensorflow.compat.v1.concat tensorflow.compat.v1.split tensorflow.compat.v1.math.cumsum tensorflow.compat.v1.variable_scope tensorflow.compat.v1.squeeze tensorflow.compat.v1.fill tensorflow.compat.v1.einsum tensorflow.compat.v1.nn.dropout
apis: tensorflow.compat.v1.not_equal tensorflow.compat.v1.zeros_initializer tensorflow.compat.v1.train.AdamOptimizer tensorflow.compat.v1.keras.backend.std tensorflow.compat.v1.layers.flatten tensorflow.compat.v1.multiply tensorflow.compat.v1.equal tensorflow.compat.v1.reduce_max tensorflow.compat.v1.train.Scaffold tensorflow.compat.v1.reduce_mean tensorflow.contrib.tpu.cross_replica_sum tensorflow.compat.v1.cast tensorflow.compat.v1.reduce_all tensorflow.compat.v1.clip_by_global_norm tensorflow.compat.v1.estimator.tpu.TPUEstimatorSpec tensorflow.compat.v1.estimator.export.PredictOutput tensorflow.compat.v1.init_from_checkpoint tensorflow.compat.v1.trainable_variables tensorflow.compat.v1.tpu.CrossShardOptimizer tensorflow.compat.v1.reduce_min tensorflow.compat.v1.train.get_global_step tensorflow.compat.v1.variable_scope numpy.sqrt tensorflow.compat.v1.train.get_or_create_global_step
apis: tensorflow.compat.v2.sysconfig.get_compile_flags tensorflow.compat.v2.sysconfig.get_link_flags
apis: matplotlib.pyplot.legend matplotlib.pyplot.savefig matplotlib.pyplot.axis matplotlib.pyplot.figure matplotlib.pyplot.xlim matplotlib.pyplot.imshow numpy.max matplotlib.pyplot.close matplotlib.pyplot.ylim matplotlib.use numpy.meshgrid numpy.linspace matplotlib.pyplot.scatter