code: # coding=utf-8
  # Copyright 2021 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.
  
  """Base class for tasks."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import abc
  import csv
  import os
  import tensorflow.compat.v1 as tf
  
  
  class Example(object):
    __metaclass__ = abc.ABCMeta
  
    def __init__(self, task_name):
      self.task_name = task_name
  
  
  class Task(object):
    """Override this class to add a new task."""
  
    __metaclass__ = abc.ABCMeta
  
    def __init__(self, config, name, long_sequences=False):
      self.config = config
      self.name = name
      self.long_sequences = long_sequences
  
    def get_examples(self, split):
      return self.load_data(split + ".tsv", split)
  
    def get_test_splits(self):
      return ["test"]
  
    def load_data(self, fname, split):
      examples = self._create_examples(
          read_tsv(os.path.join(self.config.raw_data_dir(self.name), fname),
                   max_lines=50 if self.config.debug else None),
          split)
      return examples
  
    @abc.abstractmethod
    def _create_examples(self, lines, split):
      pass
  
    @abc.abstractmethod
    def get_scorer(self):
      pass
  
    @abc.abstractmethod
    def get_feature_specs(self):
      pass
  
    @abc.abstractmethod
    def featurize(self, example, is_training):
      pass
  
    @abc.abstractmethod
    def get_prediction_module(self, bert_model, features, is_training,
                              percent_done):
      pass
  
    def __repr__(self):
      return "Task(" + self.name + ")"
  
  
  def read_tsv(input_file, quotechar=None, max_lines=None):
    """Reads a tab separated value file."""
    with tf.gfile.Open(input_file, "r") as f:
      reader = csv.reader(f, delimiter="\t", quotechar=quotechar)
      lines = []
      for i, line in enumerate(reader):
        if max_lines and i >= max_lines:
          break
        lines.append(line)
      return lines
   # coding=utf-8
  # Copyright 2021 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.
  
  """The input function of seq2act models."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from enum import Enum
  from tensor2tensor.layers import common_layers
  import tensorflow.compat.v1 as tf  # tf
  
  NUM_TOKENS_PER_OBJ = 30
  NUM_TOKENS_PER_SYN = 30
  
  
  class DataSource(Enum):
    """The class that represents word2act data source."""
    RICO_SCA = 'rico_sca'
    ANDROID_HOWTO = 'android_howto'
    PIXEL_HELP = 'pixel_help'
  
    @staticmethod
    def from_str(label):
      if label == 'rico_sca':
        return DataSource.RICO_SCA
      elif label == 'android_howto':
        return DataSource.ANDROID_HOWTO
      elif label == 'pixel_help':
        return DataSource.PIXEL_HELP
      else:
        raise ValueError('Unrecognized source %s' % label)
  
  
  MAX_UI_OBJECT_NUM = {
      DataSource.PIXEL_HELP: 93,
  }
  
  MAX_TOKEN_NUM = {
      DataSource.ANDROID_HOWTO: 30,
      DataSource.RICO_SCA: 30,
      DataSource.PIXEL_HELP: 153,
  }
  
  # ['connect_str',  token_id(connector_str)]
  # token id based on all_source_lower_case_vocab_59429
  PADDED_CONCATENATORS = [
      [5, 0, 0],
      [115, 0, 0],
      [8, 32, 0],
      [115, 8, 32],
      [12, 0, 0],
  ]
  
  CONCATENATORS_STR = [
      ', ',
      ' , ',
      ' and then ',
      ' , and then ',
      '. '
  ]
  
  
  def _construct_padding_info(data_source, load_dom_dist, load_extra):
    """Constructs the padding info tuple."""
    token_num = MAX_TOKEN_NUM[data_source]
    # Model uses this anchor padding value to mask out the padded features.
    anchor_padding_value_int = tf.cast(-1, tf.int32)
    padding_value_int = tf.cast(0, tf.int32)
    padding_value_str = tf.cast('', tf.string)
    # Tuple of (feature name, padded_shape, padded_value)
    padding_info = [
        ('task', [None], padding_value_int),
        ('rule', [], padding_value_int),
        ('verbs', [None], padding_value_int),
        ('input_refs', [None, 2], padding_value_int),
        ('obj_refs', [None, 2], padding_value_int),
        ('verb_refs', [None, 2], padding_value_int),
        ('objects', [None], padding_value_int),
        ('obj_text', [None, None, token_num], padding_value_int),
        ('obj_type', [None, None], anchor_padding_value_int),
        ('obj_clickable', [None, None], padding_value_int),
        ('obj_screen_pos', [None, None, 4], tf.cast(0, tf.int32)),
        ('obj_dom_pos', [None, None, 3], padding_value_int),
        ('agreement_count', [], padding_value_int),
        ('data_source', [], padding_value_int),
    ]
    if load_dom_dist:
      padding_info.append(('obj_dom_dist', [None, None, None], padding_value_int))
    if load_extra:
      padding_info.append(('task_id', [], padding_value_str))
      padding_info.append(('raw_task', [], padding_value_str))
      padding_info.append(('obj_raw_text', [None, None], padding_value_str))
  
    padded_shapes = {}
    padded_values = {}
    for (key, padding_shape, padding_value) in padding_info:
      padded_shapes[key] = padding_shape
      padded_values[key] = padding_value
    return padded_shapes, padded_values
  
  
  def input_fn(data_files,
               batch_size,
               repeat=-1,
               data_source=DataSource.RICO_SCA,
               required_agreement=2,
               max_range=1000,
               max_dom_pos=2000,
               max_pixel_pos=100,
               load_dom_dist=False,
               load_extra=False,
               buffer_size=8 * 1024,
               shuffle_size=8 * 1024,
               required_rule_id_list=None,
               shuffle_repeat=True,
               mean_synthetic_length=1.0,
               stddev_synthetic_length=0.0,
               load_screen=True,
               shuffle_files=True):
    """Retrieves batches of data for training.
  
    Adds padding to ensure all dimension in one batch are always same.
  
    Args:
      data_files: A list of file names to initialize the TFRecordDataset
      batch_size:  Number for the size of the batch.
      repeat: the number of times to repeat the input data.
      data_source: A DataSource instance.
      required_agreement: the minimum agreement required.
      max_range: the max range.
      max_dom_pos: the max dom pos.
      max_pixel_pos: the max screen pixels.
      load_dom_dist: whether to load the dom distance feature.
      load_extra: whether to load the raw text data.
      buffer_size: the buffer size for prefetching.
      shuffle_size: the shuffle size.
      required_rule_id_list: the list of required rule ids.
      shuffle_repeat: whether to shuffle and repeat.
      mean_synthetic_length: the mean length for synthetic sequence.
      stddev_synthetic_length: the stddev length for synthetic sequence.
      load_screen: whether to load screen features.
      shuffle_files: shuffling file names.
    Returns:
      a tf.dataset.Dateset object.
    Raises:
      ValueError: The data_format is neither 'recordio' nor 'tfrecord'.
    """
    if not isinstance(data_source, DataSource):
      assert False, 'data_source %s unsupported' % str(data_source)
    padded_shapes, padded_values = _construct_padding_info(
        data_source, load_dom_dist, load_extra)
    if not isinstance(data_files, (list,)):
      data_files = [data_files]
    all_files = tf.concat(
        values=[tf.matching_files(f) for f in data_files], axis=0)
    if repeat == -1 and shuffle_files:
      all_files = tf.random.shuffle(all_files)
    if data_files[0].endswith('.recordio'):
      dataset = tf.data.RecordIODataset(all_files)
    elif data_files[0].endswith('.tfrecord'):
      dataset = tf.data.TFRecordDataset(
          all_files, num_parallel_reads=10 if repeat == -1 else None)
    else:
      assert False, 'Data_format %s is not supported.' % data_files[0]
  
    def _map_fn(x):
      return parse_tf_example(x, data_source, max_range, max_dom_pos,
                              max_pixel_pos, load_dom_dist=load_dom_dist,
                              load_extra=load_extra,
                              append_eos=(data_source != DataSource.RICO_SCA or
                                          mean_synthetic_length == 1.0),
                              load_screen=load_screen)
    dataset = dataset.map(_map_fn)
    def _is_enough_agreement(example):
      return tf.greater_equal(example['agreement_count'], required_agreement)
    dataset = dataset.filter(_is_enough_agreement)
  
    def _length_filter(example):
      return tf.less(tf.shape(example['obj_refs'])[0], 20)
    dataset = dataset.filter(_length_filter)
  
    def _filter_data_by_rule(example, rule_id_list):
      return tf.reduce_any(
          [tf.equal(example['rule'], rule_id) for rule_id in rule_id_list])
    if data_source == DataSource.RICO_SCA and required_rule_id_list is not None:
      dataset = dataset.filter(
          lambda x: _filter_data_by_rule(x, required_rule_id_list))
  
    # (TODO: liyang) tf.data.experimental.bucket_by_sequence_length
    if shuffle_repeat:
      dataset = dataset.apply(tf.data.experimental.shuffle_and_repeat(
          shuffle_size, count=repeat))
    dataset = dataset.padded_batch(
        batch_size, padded_shapes=padded_shapes, padding_values=padded_values)
    if data_source == DataSource.RICO_SCA and mean_synthetic_length > 1.0:
      def _stitch_fn(x):
        return _batch_stitch(x, mean_length=mean_synthetic_length,
                             stddev=stddev_synthetic_length)
      dataset = dataset.map(_stitch_fn)
    dataset = dataset.prefetch(buffer_size=buffer_size)
    return dataset
  
  
  def hybrid_input_fn(data_files_list,
                      data_source_list,
                      batch_size_list,
                      max_range=1000,
                      max_dom_pos=2000,
                      max_pixel_pos=100,
                      load_dom_dist=False,
                      load_extra=False,
                      buffer_size=8 * 1024,
                      mean_synthetic_length=1.0,
                      stddev_synthetic_length=0.0,
                      hybrid_batch_size=128,
                      boost_input=False,
                      load_screen=True,
                      shuffle_size=1024):
    """Combines multiple datasouces."""
    mixed_dataset = None
    for data_files, data_source, batch_size in zip(
        data_files_list, data_source_list, batch_size_list):
      dataset = input_fn(data_files, batch_size, repeat=-1,
                         data_source=data_source,
                         required_agreement=-1,
                         max_range=max_range, max_dom_pos=max_dom_pos,
                         max_pixel_pos=max_pixel_pos,
                         load_dom_dist=load_dom_dist,
                         load_extra=load_extra,
                         buffer_size=0,
                         mean_synthetic_length=mean_synthetic_length,
                         stddev_synthetic_length=stddev_synthetic_length,
                         shuffle_repeat=False,
                         load_screen=load_screen)
      if mixed_dataset is None:
        mixed_dataset = dataset
      else:
        mixed_dataset = dataset.concatenate(mixed_dataset)
  
    mixed_dataset = mixed_dataset.unbatch()
    # Boost input examples
    if boost_input:
      def _input_booster(example):
        with tf.control_dependencies([tf.rank(example['input_refs']), 2]):
          has_input = tf.reduce_any(
              tf.greater(example['input_refs'][:, 1],
                         example['input_refs'][:, 0]))
          return tf.logical_or(has_input, tf.less(tf.random_uniform([]), 0.1))
      dataset = dataset.filter(_input_booster)
    # Remix single examples
    mixed_dataset = mixed_dataset.shuffle(hybrid_batch_size * shuffle_size)
    # Batch again
    padded_shapes, padded_values = _construct_padding_info(
        data_source_list[0], load_dom_dist, load_extra)
    mixed_dataset = mixed_dataset.padded_batch(
        hybrid_batch_size, padded_shapes=padded_shapes,
        padding_values=padded_values)
    mixed_dataset = mixed_dataset.repeat()
    mixed_dataset = mixed_dataset.prefetch(buffer_size=buffer_size)
    return mixed_dataset
  
  
  def parse_tf_example(example_proto,
                       data_source,
                       max_range=100,
                       max_dom_pos=2000,
                       max_pixel_pos=100,
                       load_dom_dist=False,
                       load_extra=False,
                       append_eos=True,
                       load_screen=True):
    """Parses an example TFRecord proto into dictionary of tensors.
  
    Args:
      example_proto: TFRecord format proto that contains screen information.
      data_source: A DataSource instance.
      max_range: the max range.
      max_dom_pos: the maximum dom positoin.
      max_pixel_pos: the max dom position.
      load_dom_dist: whether to load the feature.
      load_extra: whether to load the extra data for debugging.
      append_eos: whether to append eos.
      load_screen: whether to load screen features.
    Returns:
      feature: The parsed tensor dictionary with the input feature data
      label: The parsed label tensor with the input label for the feature
    """
    feature_spec = {
        'instruction_word_id_seq':
            tf.FixedLenSequenceFeature([], tf.int64, allow_missing=True),
        'input_str_position_seq':
            tf.FixedLenSequenceFeature([], tf.int64, allow_missing=True),
        'obj_desc_position_seq':
            tf.FixedLenSequenceFeature([], tf.int64, allow_missing=True),
        'verb_str_position_seq':
            tf.FixedLenSequenceFeature([], tf.int64, allow_missing=True),
        'agreement_count':
            tf.FixedLenSequenceFeature([], tf.int64, allow_missing=True),
        'instruction_rule_id':
            tf.FixedLenSequenceFeature([], tf.int64, allow_missing=True)
    }
    if load_screen:
      feature_spec['verb_id_seq'] = tf.FixedLenSequenceFeature(
          [], tf.int64, allow_missing=True)
      feature_spec['ui_target_id_seq'] = tf.FixedLenSequenceFeature(
          [], tf.int64, allow_missing=True)
      feature_spec['ui_obj_word_id_seq'] = tf.FixedLenSequenceFeature(
          [], tf.int64, allow_missing=True)
      feature_spec['ui_obj_type_id_seq'] = tf.FixedLenSequenceFeature(
          [], tf.int64, allow_missing=True)
      feature_spec['ui_obj_clickable_seq'] = tf.FixedLenSequenceFeature(
          [], tf.int64, allow_missing=True)
      feature_spec['ui_obj_cord_x_seq'] = tf.FixedLenSequenceFeature(
          [], tf.float32, allow_missing=True)
      feature_spec['ui_obj_cord_y_seq'] = tf.FixedLenSequenceFeature(
          [], tf.float32, allow_missing=True)
      feature_spec['ui_obj_dom_location_seq'] = tf.FixedLenSequenceFeature(
          [], tf.int64, allow_missing=True)
  
    if load_dom_dist:
      feature_spec['ui_obj_dom_distance'] = tf.FixedLenSequenceFeature(
          [], tf.int64, allow_missing=True)
    if load_extra:
      feature_spec['instruction_str'] = tf.FixedLenSequenceFeature(
          [], tf.string, allow_missing=True)
      feature_spec['task_id'] = tf.FixedLenSequenceFeature(
          [], tf.string, allow_missing=True)
      feature_spec['ui_obj_str_seq'] = tf.FixedLenSequenceFeature(
          [], tf.string, allow_missing=True)
  
    feature_dict = tf.parse_single_example(example_proto, feature_spec)
  
    for key in feature_dict:
      if feature_dict[key].dtype == tf.int64:
        feature_dict[key] = tf.cast(feature_dict[key], tf.int32)
    if data_source == DataSource.ANDROID_HOWTO:
      tf.logging.info('Parsing android_howto dataset')
      feature = _process_android_howto(feature_dict, max_range=max_range,
                                       load_dom_dist=load_dom_dist,
                                       load_extra=load_extra)
    elif data_source == DataSource.RICO_SCA:
      tf.logging.info('Parsing synthetic dataset')
      feature = _process_rico_sca(
          feature_dict, max_range=max_range, max_dom_pos=max_dom_pos,
          load_dom_dist=load_dom_dist,
          load_extra=load_extra,
          load_screen=load_screen)
    elif data_source == DataSource.PIXEL_HELP:
      tf.logging.info('Parsing test dataset')
      feature = _process_pixel_help(feature_dict, data_source,
                                    load_dom_dist=load_dom_dist,
                                    load_extra=load_extra)
    else:
      raise ValueError('Unsupported datasource %s' % str(data_source))
    # Remove padding from "task"
    feature['task'] = tf.boolean_mask(feature['task'],
                                      tf.not_equal(feature['task'], 0))
    feature['obj_screen_pos'] = tf.to_int32(
        feature['obj_screen_pos'] * (max_pixel_pos - 1))
    # Appending EOS and padding to match the appended length
    if append_eos:
      feature['input_refs'] = tf.pad(feature['input_refs'], [[0, 1], [0, 0]])
      feature['obj_refs'] = tf.pad(feature['obj_refs'], [[0, 1], [0, 0]])
      step_num = tf.size(feature['task'])
      feature['verb_refs'] = tf.concat(
          [feature['verb_refs'], [[step_num, step_num + 1]]], axis=0)
      feature['task'] = tf.pad(feature['task'], [[0, 1]], constant_values=1)
      feature['obj_text'] = tf.pad(feature['obj_text'], [[0, 1], [0, 0], [0, 0]])
      feature['obj_clickable'] = tf.pad(feature['obj_clickable'],
                                        [[0, 1], [0, 0]])
      feature['obj_type'] = tf.pad(
          feature['obj_type'], [[0, 1], [0, 0]], constant_values=-1)
      feature['obj_screen_pos'] = tf.pad(feature['obj_screen_pos'],
                                         [[0, 1], [0, 0], [0, 0]])
      feature['obj_dom_pos'] = tf.pad(feature['obj_dom_pos'],
                                      [[0, 1], [0, 0], [0, 0]])
      if load_dom_dist:
        feature['obj_dom_dist'] = tf.pad(feature['obj_dom_dist'],
                                         [[0, 1], [0, 0], [0, 0]])
      feature['objects'] = tf.pad(feature['objects'], [[0, 1]])
      feature['verbs'] = tf.pad(feature['verbs'], [[0, 1]])
    return feature
  
  
  def _bound_refs(feature, max_range):
    """Makes sure the refs are in the allowed range."""
    for key in feature:
      if not key.endswith('_refs'):
        continue
      feature[key] = tf.where(
          tf.greater(feature[key][:, 1] - feature[key][:, 0], max_range),
          tf.stack([feature[key][:, 0], feature[key][:, 0] + max_range], axis=1),
          feature[key])
  
  
  def _process_android_howto(feature_dict, max_range, load_dom_dist=False,
                             load_extra=False):
    """Processes webanswer feature dictionary.
  
    Args:
      feature_dict: feature dictionary
      max_range: the max range.
      load_dom_dist: whether to load the dom distance feature.
      load_extra: whether to load the extra data for debugging.
    Returns:
      A processed feature dictionary.
    """
  
    feature = {
        'task': tf.reshape(feature_dict['instruction_word_id_seq'], [-1]),
        'input_refs': tf.reshape(feature_dict['input_str_position_seq'], [-1, 2]),
        'obj_refs': tf.reshape(feature_dict['obj_desc_position_seq'], [-1, 2]),
        'verb_refs': tf.reshape(feature_dict['verb_str_position_seq'], [-1, 2]),
        'agreement_count': tf.reshape(feature_dict['agreement_count'], [])
    }
    if load_extra:
      feature['task_id'] = tf.constant('empty_task_id', dtype=tf.string)
      feature['raw_task'] = tf.reshape(feature_dict['instruction_str'], [])
    _bound_refs(feature, max_range)
    _load_fake_screen(feature, load_extra, load_dom_dist)
    return feature
  
  
  def _load_fake_screen(feature, load_extra, load_dom_dist):
    """Loads a fake screen."""
    # Fills in fake ui object features into feature dictionary.
    step_num = tf.shape(feature['verb_refs'])[0]
    obj_num = 1
    if load_extra:
      feature['obj_raw_text'] = tf.fill([step_num, obj_num], '')
    feature['data_source'] = tf.constant(1, dtype=tf.int32)
    feature['obj_text'] = tf.zeros([step_num, obj_num, NUM_TOKENS_PER_OBJ],
                                   tf.int32)
    feature['obj_type'] = tf.cast(tf.fill([step_num, obj_num], -1), tf.int32)
    feature['obj_clickable'] = tf.zeros([step_num, obj_num], tf.int32)
    feature['obj_screen_pos'] = tf.zeros([step_num, obj_num, 4], tf.float32)
    feature['obj_dom_pos'] = tf.zeros([step_num, obj_num, 3], tf.int32)
    if load_dom_dist:
      feature['obj_dom_dist'] = tf.zeros([step_num, obj_num, obj_num], tf.int32)
    feature['objects'] = tf.zeros([step_num], tf.int32)
    feature['verbs'] = tf.zeros([step_num], tf.int32)
    feature['rule'] = tf.constant(5, dtype=tf.int32)
  
  
  def _batch_stitch(features, mean_length=4.0, stddev=2.0):
    """Stitches a batch of single-step data to a batch of multi-step data."""
    batch_size = common_layers.shape_list(features['task'])[0]
    num_sequences = tf.maximum(
        tf.to_int32(tf.to_float(batch_size) / mean_length), 1)
    lengths = tf.random.truncated_normal(shape=[num_sequences],
                                         mean=mean_length, stddev=stddev)
    max_length = tf.reduce_max(lengths) * (
        tf.to_float(batch_size) / tf.reduce_sum(lengths))
    max_length = tf.to_int32(tf.ceil(max_length))
    total_items = max_length * num_sequences
    num_paddings = total_items - batch_size
    indices = tf.random.shuffle(tf.range(total_items))
    for key in features:
      shape_list = common_layers.shape_list(features[key])
      assert len(shape_list) >= 1
      with tf.control_dependencies([
          tf.assert_greater_equal(num_paddings, 0,
                                  name='num_paddings_positive')]):
        paddings = [[0, num_paddings]] + [[0, 0]] * (len(shape_list) - 1)
      features[key] = tf.pad(features[key], paddings,
                             constant_values=-1 if key == 'obj_type' else 0)
      features[key] = tf.gather(features[key], indices)
      shape = [num_sequences, max_length]
      if len(shape_list) >= 2:
        shape += shape_list[1:]
      features[key] = tf.reshape(features[key], shape)
    # Remove all-padding seqs
    step_mask = tf.reduce_any(tf.greater(features['task'], 1), axis=-1)
    mask = tf.reduce_any(step_mask, axis=-1)
    step_mask = tf.boolean_mask(step_mask, mask)
    for key in features:
      features[key] = tf.boolean_mask(features[key], mask=mask)
    num_sequences = tf.shape(features['task'])[0]
    # Sort steps within each seq
    _, step_indices = tf.math.top_k(tf.to_int32(step_mask), k=max_length)
    step_indices = step_indices + tf.expand_dims(
        tf.range(num_sequences) * max_length, 1)
    step_indices = tf.reshape(step_indices, [-1])
    for key in features:
      shape_list = common_layers.shape_list(features[key])
      features[key] = tf.gather(tf.reshape(features[key], [-1] + shape_list[2:]),
                                step_indices)
      features[key] = tf.reshape(features[key], shape_list)
    features = _stitch(features)
    return features
  
  
  def _stitch(features):
    """Stitch features on the first dimension."""
    full_mask = tf.greater(features['task'], 1)
    step_mask = tf.reduce_any(full_mask, axis=-1)
    step_mask_exclude_last = tf.pad(step_mask,
                                    [[0, 0], [0, 1]],
                                    constant_values=False)[:, 1:]
    num_sequences = common_layers.shape_list(features['task'])[0]
    num_steps = common_layers.shape_list(features['task'])[1]
    connectors = tf.constant(PADDED_CONCATENATORS)
    # Select connectors
    connector_indices = tf.random.uniform(
        [num_sequences * num_steps], minval=0,
        maxval=len(PADDED_CONCATENATORS), dtype=tf.int32)
    selected_connectors = tf.reshape(
        tf.gather(connectors, connector_indices),
        [num_sequences, num_steps, len(PADDED_CONCATENATORS[0])])
    selected_connectors = tf.multiply(
        selected_connectors,
        tf.expand_dims(tf.to_int32(step_mask_exclude_last), 2),
        name='connector_mask')
    features['task'] = tf.concat([features['task'], selected_connectors], axis=-1)
    ref_offsets = tf.expand_dims(
        tf.cumsum(tf.reduce_sum(tf.to_int32(tf.greater(features['task'], 1)), -1),
                  exclusive=True, axis=-1), 2)
    features['task'] = tf.reshape(features['task'], [num_sequences, -1])
    full_mask = tf.greater(features['task'], 1)
    full_mask_int = tf.to_int32(full_mask)
    indices = tf.where(tf.sequence_mask(lengths=tf.reduce_sum(full_mask_int, -1)))
    values = tf.boolean_mask(tf.reshape(features['task'], [-1]),
                             tf.reshape(full_mask, [-1]))
    sparse_task = tf.sparse.SparseTensor(
        indices=indices, values=values,
        dense_shape=tf.to_int64(tf.shape(features['task'])))
    # Stitch task and raw_task
    stitched_features = {}
    stitched_features['task'] = tf.sparse_tensor_to_dense(sparse_task)
    max_len = tf.reduce_max(
        tf.reduce_sum(tf.to_int32(tf.greater(stitched_features['task'], 1)), -1))
    stitched_features['task'] = stitched_features['task'][:, :max_len]
    if 'raw_task' in features:
      connector_strs = tf.reshape(
          tf.gather(tf.constant(CONCATENATORS_STR), connector_indices),
          [num_sequences, num_steps])
      masked_connector_strs = tf.where(
          step_mask_exclude_last,
          connector_strs, tf.fill(tf.shape(connector_strs), ''))
      stitched_features['raw_task'] = tf.strings.reduce_join(
          tf.strings.reduce_join(tf.concat([
              tf.expand_dims(features['raw_task'], 2),
              tf.expand_dims(masked_connector_strs, 2)], axis=2), axis=-1), -1)
    # Stitch screen sequences
    action_lengths = tf.reduce_sum(tf.to_int32(
        tf.greater(features['verb_refs'][:, :, 0, 1],
                   features['verb_refs'][:, :, 0, 0])), -1)
    max_action_length = tf.reduce_max(action_lengths)
    def _pad(tensor, padding_value=0):
      shape_list = common_layers.shape_list(tensor)
      assert len(shape_list) >= 2
      padding_list = [[0, 0], [0, 1]] + [[0, 0]] * (len(shape_list) - 2)
      return tf.pad(tensor[:, :max_action_length],
                    padding_list, constant_values=padding_value)
    for key in features.keys():
      if key.endswith('_refs'):
        features[key] = tf.squeeze(features[key], 2)
        ref_mask = tf.expand_dims(tf.to_int32(
            tf.not_equal(features[key][:, :, 0],
                         features[key][:, :, 1])), 2)
        stitched_features[key] = tf.multiply(
            (features[key] + ref_offsets), ref_mask, name='ref_mask')
        stitched_features[key] = _pad(stitched_features[key])
      elif key in ['verbs', 'objects', 'consumed', 'obj_dom_pos',
                   'obj_text', 'obj_type', 'obj_clickable', 'obj_screen_pos',
                   'verb_refs', 'obj_refs', 'input_refs', 'obj_dom_dist']:
        features[key] = tf.squeeze(features[key], 2)
        stitched_features[key] = features[key]
        stitched_features[key] = _pad(
            stitched_features[key],
            padding_value=-1 if key == 'obj_type' else 0)
      elif key not in ['task', 'raw_task']:
        stitched_features[key] = features[key][:, 0]
    # Append eos to 'task'
    stitched_features['task'] = tf.pad(stitched_features['task'],
                                       [[0, 0], [0, 1]])
    task_mask = tf.to_int32(tf.greater(stitched_features['task'], 1))
    task_eos_mask = tf.pad(task_mask, [[0, 0], [1, 0]], constant_values=1)[:, :-1]
    stitched_features['task'] = stitched_features['task'] + (
        task_eos_mask - task_mask)
    # Append eos
    verb_mask = tf.to_int32(tf.greater(stitched_features['verbs'], 1))
    verb_eos_mask = tf.pad(verb_mask, [[0, 0], [1, 0]], constant_values=1)[:, :-1]
    verb_eos = verb_eos_mask - verb_mask
    stitched_features['verbs'] = stitched_features['verbs'] + verb_eos
    # Append last step refs to 'verb_refs'
    task_lengths = tf.where(tf.equal(stitched_features['task'], 1))[:, 1]
    eos_pos = tf.to_int32(tf.stack([task_lengths, task_lengths + 1], axis=1))
    action_mask = tf.to_int32(
        tf.sequence_mask(action_lengths, max_action_length + 1))
    action_and_eos_mask = tf.pad(action_mask, [[0, 0], [1, 0]],
                                 constant_values=1)[:, :-1]
    verb_ref_eos = action_and_eos_mask - action_mask
    eos_refs = tf.multiply(
        tf.tile(tf.expand_dims(eos_pos, 1), [1, max_action_length + 1, 1]),
        tf.expand_dims(verb_ref_eos, 2), name='verb_ref_eos')
    stitched_features['verb_refs'] += eos_refs
    return stitched_features
  
  
  def _process_rico_sca(feature_dict, max_range, max_dom_pos,
                        load_dom_dist=False, load_extra=False, load_screen=True):
    """Processes one_shot feature dictionary.
  
    Args:
      feature_dict: feature dictionary
      max_range: the max range.
      max_dom_pos: the max dom pos.
      load_dom_dist: whether to load the dom distance feature.
      load_extra: whether to load the extra data for debugging.
      load_screen: whether to load the screen features.
    Returns:
      A processed feature dictionary.
    """
    phrase_count = tf.size(feature_dict['obj_desc_position_seq']) // 2
    feature = {
        'task':
            tf.reshape(feature_dict['instruction_word_id_seq'],
                       [phrase_count, NUM_TOKENS_PER_SYN]),
        'input_refs':
            tf.reshape(feature_dict['input_str_position_seq'],
                       [phrase_count, 1, 2]),
        'obj_refs':
            tf.reshape(feature_dict['obj_desc_position_seq'],
                       [phrase_count, 1, 2]),
        'verb_refs':
            tf.reshape(feature_dict['verb_str_position_seq'],
                       [phrase_count, 1, 2]),
        'rule':
            tf.reshape(feature_dict['instruction_rule_id'], [phrase_count]),
    }
    selected_synthetic_action_idx = tf.random_uniform(
        shape=(), minval=0, maxval=phrase_count, dtype=tf.int32)
    for key in feature:
      feature[key] = feature[key][selected_synthetic_action_idx]
    if load_extra:
      feature['raw_task'] = tf.reshape(
          feature_dict['instruction_str'],
          [phrase_count])[selected_synthetic_action_idx]
      feature['task_id'] = tf.constant('empty_task_id', dtype=tf.string)
    if load_screen:
      feature['verbs'] = tf.reshape(
          feature_dict['verb_id_seq'],
          [phrase_count, 1])[selected_synthetic_action_idx]
      feature['objects'] = tf.reshape(
          feature_dict['ui_target_id_seq'],
          [phrase_count, 1])[selected_synthetic_action_idx]
      feature['obj_text'] = tf.reshape(feature_dict['ui_obj_word_id_seq'],
                                       [1, -1, NUM_TOKENS_PER_OBJ])
      feature['obj_type'] = tf.reshape(
          feature_dict['ui_obj_type_id_seq'], [1, -1])
      feature['obj_clickable'] = tf.reshape(feature_dict['ui_obj_clickable_seq'],
                                            [1, -1])
      def _make_obj_screen_pos():
        return tf.concat([
            tf.reshape(feature_dict['ui_obj_cord_x_seq'], [1, -1, 2]),
            tf.reshape(feature_dict['ui_obj_cord_y_seq'], [1, -1, 2])
        ], 2)
  
      feature['obj_screen_pos'] = tf.cond(
          tf.equal(
              tf.size(feature_dict['ui_obj_cord_x_seq']),
              0), lambda: tf.fill([1, tf.shape(feature['obj_type'])[1], 4], 0.),
          _make_obj_screen_pos)
      feature['obj_dom_pos'] = tf.reshape(feature_dict['ui_obj_dom_location_seq'],
                                          [1, -1, 3])
      feature['obj_dom_pos'] = tf.minimum(feature['obj_dom_pos'], max_dom_pos - 1)
      if load_dom_dist:
        num_ui_obj = tf.to_int32(
            tf.sqrt(tf.to_float(tf.size(feature_dict['ui_obj_dom_distance']))))
        feature['obj_dom_dist'] = tf.reshape(feature_dict['ui_obj_dom_distance'],
                                             [1, num_ui_obj, num_ui_obj])
      if load_extra:
        feature['obj_raw_text'] = tf.reshape(feature_dict['ui_obj_str_seq'],
                                             [1, -1])
    else:
      _load_fake_screen(feature, load_extra, load_dom_dist)
    _bound_refs(feature, max_range)
    feature['data_source'] = tf.constant(0, dtype=tf.int32)
    feature['agreement_count'] = tf.constant(100, dtype=tf.int32)
  
    return feature
  
  
  def _process_pixel_help(feature_dict, data_source, load_dom_dist=False,
                          load_extra=False):
    """Processes testing data feature dictionary.
  
    Args:
      feature_dict: feature dictionary
      data_source: TEST_PIXEL_HELP
      load_dom_dist: whether to load the dom distance feature.
      load_extra: whether to load the extra data for debugging.
    Returns:
      A processed feature dictionary.
    """
    step_num = tf.size(feature_dict['verb_id_seq'])
    feature = {
        'task':
            tf.reshape(feature_dict['instruction_word_id_seq'], [-1]),
        'obj_text':
            tf.reshape(feature_dict['ui_obj_word_id_seq'], [
                step_num, MAX_UI_OBJECT_NUM[data_source],
                MAX_TOKEN_NUM[data_source]
            ]),
        'obj_type':
            tf.reshape(feature_dict['ui_obj_type_id_seq'],
                       [step_num, MAX_UI_OBJECT_NUM[data_source]]),
        'obj_clickable':
            tf.reshape(feature_dict['ui_obj_clickable_seq'],
                       [step_num, MAX_UI_OBJECT_NUM[data_source]]),
        # pylint: disable=g-long-ternary
        'obj_screen_pos': (
            tf.reshape(tf.concat([
                tf.reshape(feature_dict['ui_obj_cord_x_seq'], [step_num, -1, 2]),
                tf.reshape(feature_dict['ui_obj_cord_y_seq'], [step_num, -1, 2])
            ], axis=2), [step_num, MAX_UI_OBJECT_NUM[data_source], 4])),
        'obj_dom_pos':
            tf.reshape(feature_dict['ui_obj_dom_location_seq'],
                       [step_num, MAX_UI_OBJECT_NUM[data_source], 3]),
        'verbs':
            tf.reshape(feature_dict['verb_id_seq'], [step_num]),
        'objects':
            tf.reshape(feature_dict['ui_target_id_seq'], [step_num]),
        'input_refs':
            tf.reshape(feature_dict['input_str_position_seq'], [step_num, 2]),
        'obj_refs':
            tf.reshape(feature_dict['obj_desc_position_seq'], [step_num, 2]),
        'verb_refs':  # No data for Pixel on the field
            tf.zeros([step_num, 2], tf.int32),
        'agreement_count':
            tf.constant(100, dtype=tf.int32),
    }
    if load_dom_dist:
      feature['obj_dom_dist'] = tf.reshape(
          feature_dict['ui_obj_dom_distance'],
          [step_num, MAX_UI_OBJECT_NUM[data_source],
           MAX_UI_OBJECT_NUM[data_source]])
    feature['rule'] = tf.constant(5, dtype=tf.int32)
    if load_extra:
      feature['task_id'] = tf.reshape(feature_dict['task_id'], [])
      feature['raw_task'] = tf.reshape(feature_dict['instruction_str'], [])
      feature['obj_raw_text'] = tf.reshape(
          feature_dict['ui_obj_str_seq'],
          [step_num, MAX_UI_OBJECT_NUM[data_source]])
    feature['data_source'] = tf.constant(2, dtype=tf.int32)
    return feature
   # coding=utf-8
  # Copyright 2021 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.
  
  # Lint as: python3
  r"""Approximate the distribution's partition function with a spline.
  
  This script generates values for the distribution's partition function and then
  fits a cubic hermite spline to those values, which is then stored to disk.
  To run this script, assuming you're in this directory, run:
    python -m robust_loss.fit_partition_spline_test
  This script will likely never have to be run again, and is provided here for
  completeness and reproducibility, or in case someone decides to modify
  distribution.partition_spline_curve() in the future in case they find a better
  curve. If the user wants a more accurate spline approximation, this can be
  obtained by modifying the `x_max`, `x_scale`, and `redundancy` parameters in the
  code below, but this should only be done with care.
  """
  
  from absl import app
  import numpy as np
  import tensorflow.compat.v2 as tf
  from robust_loss import cubic_spline
  from robust_loss import distribution
  from robust_loss import general
  
  tf.enable_v2_behavior()
  
  
  def numerical_base_partition_function(alpha):
    """Numerically approximate the partition function Z(alpha)."""
    # Generate values `num_samples` values in [-x_max, x_max], with more samples
    # near the origin as `power` is set to larger values.
    num_samples = 2**24 + 1  # We want an odd value so that 0 gets sampled.
    x_max = 10**10
    power = 6
    t = t = tf.linspace(
        tf.constant(-1, tf.float64), tf.constant(1, tf.float64), num_samples)
    t = tf.sign(t) * tf.abs(t)**power
    x = t * x_max
  
    # Compute losses for the values, then exponentiate the negative losses and
    # integrate with the trapezoid rule to get the partition function.
    losses = general.lossfun(x, alpha, np.float64(1))
    y = tf.math.exp(-losses)
    partition = tf.reduce_sum((y[1:] + y[:-1]) * (x[1:] - x[:-1])) / 2.
    return partition
  
  
  def main(argv):
    if len(argv) > 1:
      raise app.UsageError('Too many command-line arguments.')
  
    # Parameters governing how the x coordinate of the spline will be laid out.
    # We will construct a spline with knots at
    #   [0 : 1 / x_scale : x_max],
    # by fitting it to values sampled at
    #   [0 : 1 / (x_scale * redundancy) : x_max]
    x_max = 12
    x_scale = 1024
    redundancy = 4  # Must be >= 2 for the spline to be useful.
  
    spline_spacing = 1. / (x_scale * redundancy)
    x_knots = np.arange(
        0, x_max + spline_spacing, spline_spacing, dtype=np.float64)
    table = []
    # We iterate over knots, and for each knot recover the alpha value
    # corresponding to that knot with inv_partition_spline_curve(), and then
    # with that alpha we accurately approximate its partition function using
    # numerical_base_partition_function().
    for x_knot in x_knots:
      alpha = distribution.inv_partition_spline_curve(x_knot).numpy()
      partition = numerical_base_partition_function(alpha).numpy()
      table.append((x_knot, alpha, partition))
      print(table[-1])
  
    table = np.array(table)
    x = table[:, 0]
    alpha = table[:, 1]
    y_gt = np.log(table[:, 2])
  
    # We grab the values from the true log-partition table that correpond to
    # knots, by looking for where x * x_scale is an integer.
    mask = np.abs(np.round(x * x_scale) - (x * x_scale)) <= 1e-8
    values = y_gt[mask]
  
    # Initialize `tangents` using a central differencing scheme.
    values_pad = np.concatenate([[values[0] - values[1] + values[0]], values,
                                 [values[-1] - values[-2] + values[-1]]], 0)
    tangents = (values_pad[2:] - values_pad[:-2]) / 2.
  
    # Construct the spline's value and tangent TF variables, constraining the last
    # knot to have a fixed value Z(infinity) and a tangent of zero.
    n = len(values)
    tangents = tf.Variable(tangents, tf.float64)
    values = tf.Variable(values, tf.float64)
  
    # Fit the spline.
    num_iters = 10001
  
    optimizer = tf.keras.optimizers.SGD(learning_rate=1e-9, momentum=0.99)
  
    trace = []
    for ii in range(num_iters):
      with tf.GradientTape() as tape:
        tape.watch([values, tangents])
        # Fix the endpoint to be a known constant with a zero tangent.
        i_values = tf.where(
            np.arange(n) == (n - 1),
            tf.ones_like(values) * 0.70526025442689566, values)
        i_tangents = tf.where(
            np.arange(n) == (n - 1), tf.zeros_like(tangents), tangents)
        i_y = cubic_spline.interpolate1d(x * x_scale, i_values, i_tangents)
        # We minimize the maximum residual, which makes for a very ugly
        # optimization problem but works well in practice.
        i_loss = tf.reduce_max(tf.abs(i_y - y_gt))
        grads = tape.gradient(i_loss, [values, tangents])
        optimizer.apply_gradients(zip(grads, [values, tangents]))
      trace.append(i_loss.numpy())
      if (ii % 200) == 0:
        print('{:5d}: {:e}'.format(ii, trace[-1]))
  
    mask = alpha <= 4
    max_error_a4 = np.max(np.abs(i_y[mask] - y_gt[mask]))
    max_error = np.max(np.abs(i_y - y_gt))
    print('Max Error (a <= 4): {:e}'.format(max_error_a4))
    print('Max Error: {:e}'.format(max_error))
  
    # Just a sanity-check on the error.
    assert max_error_a4 <= 5e-7
    assert max_error <= 5e-7
  
    # Save the spline to disk.
    np.savez(
        './data/partition_spline.npz',
        x_scale=x_scale,
        values=i_values.numpy(),
        tangents=i_tangents.numpy())
  
  
  if __name__ == '__main__':
    app.run(main)
   # coding=utf-8
  # Copyright 2021 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.
  
  r"""TCN loss for unsupervised training."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import tensorflow.compat.v2 as tf
  
  from tcc.algos.algorithm import Algorithm
  from tcc.config import CONFIG
  from tcc.utils import get_cnn_feats
  from tcc.utils import set_learning_phase
  
  
  def _npairs_loss(labels, embeddings_anchor, embeddings_positive, reg_lambda):
    """Returns n-pairs metric loss."""
    reg_anchor = tf.reduce_mean(tf.reduce_sum(tf.square(embeddings_anchor), 1))
    reg_positive = tf.reduce_mean(tf.reduce_sum(
        tf.square(embeddings_positive), 1))
    l2loss = 0.25 * reg_lambda * (reg_anchor + reg_positive)
  
    # Get per pair similarities.
    similarity_matrix = tf.matmul(
        embeddings_anchor, embeddings_positive, transpose_a=False,
        transpose_b=True)
  
    # Reshape [batch_size] label tensor to a [batch_size, 1] label tensor.
    lshape = tf.shape(labels)
    assert lshape.shape == 1
    labels = tf.reshape(labels, [lshape[0], 1])
  
    labels_remapped = tf.cast(
        tf.equal(labels, tf.transpose(labels)), tf.float32)
    labels_remapped /= tf.reduce_sum(labels_remapped, 1, keepdims=True)
  
    # Add the softmax loss.
    xent_loss = tf.nn.softmax_cross_entropy_with_logits(
        logits=similarity_matrix, labels=labels_remapped)
    xent_loss = tf.reduce_mean(xent_loss)
  
    return l2loss + xent_loss
  
  
  def single_sequence_loss(embs, num_steps):
    """Returns n-pairs loss for a single sequence."""
  
    labels = tf.range(num_steps)
    labels = tf.stop_gradient(labels)
    embeddings_anchor = embs[0::2]
    embeddings_positive = embs[1::2]
    loss = _npairs_loss(labels, embeddings_anchor, embeddings_positive,
                        reg_lambda=CONFIG.TCN.REG_LAMBDA)
    return loss
  
  
  class TCN(Algorithm):
    """Time-contrastive Network."""
  
    @set_learning_phase
    def call(self, data, steps, seq_lens, training):
      """One pass through the model."""
      cnn = self.model['cnn']
      emb = self.model['emb']
  
      if training:
        num_steps = CONFIG.TRAIN.NUM_FRAMES * CONFIG.DATA.NUM_STEPS
      else:
        num_steps = CONFIG.EVAL.NUM_FRAMES * CONFIG.DATA.NUM_STEPS
  
      # Number of steps is doubled due to sampling of positives and anchors.
      cnn_feats = get_cnn_feats(cnn, data, training, 2 * num_steps)
  
      if training:
        num_steps = CONFIG.TRAIN.NUM_FRAMES
      else:
        num_steps = CONFIG.EVAL.NUM_FRAMES
  
      embs = emb(cnn_feats, 2 * num_steps)
      embs = tf.stack(tf.split(embs, 2 * num_steps, axis=0), axis=1)
  
      return embs
  
    def compute_loss(self, embs, steps, seq_lens, global_step, training,
                     frame_labels, seq_labels):
      if training:
        num_steps = CONFIG.TRAIN.NUM_FRAMES
        batch_size = CONFIG.TRAIN.BATCH_SIZE
      else:
        num_steps = CONFIG.EVAL.NUM_FRAMES
        batch_size = CONFIG.EVAL.BATCH_SIZE
      losses = []
      for i in xrange(batch_size):
        losses.append(single_sequence_loss(embs[i], num_steps))
      loss = tf.reduce_mean(tf.stack(losses))
      return loss
  
apis: tensorflow.compat.v1.gfile.Open
apis: tensorflow.compat.v1.not_equal tensorflow.compat.v1.parse_single_example tensorflow.compat.v1.data.RecordIODataset tensorflow.compat.v1.expand_dims tensorflow.compat.v1.reduce_sum tensorflow.compat.v1.shape tensorflow.compat.v1.multiply tensorflow.compat.v1.equal tensorflow.compat.v1.reduce_max tensorflow.compat.v1.greater tensorflow.compat.v1.boolean_mask tensorflow.compat.v1.sparse_tensor_to_dense tensorflow.compat.v1.gather tensorflow.compat.v1.constant tensorflow.compat.v1.random.shuffle tensorflow.compat.v1.cast tensorflow.compat.v1.data.TFRecordDataset tensorflow.compat.v1.random.truncated_normal tensorflow.compat.v1.stack tensorflow.compat.v1.to_int32 tensorflow.compat.v1.pad tensorflow.compat.v1.reshape tensorflow.compat.v1.reduce_any tensorflow.compat.v1.sequence_mask tensorflow.compat.v1.minimum tensorflow.compat.v1.concat tensorflow.compat.v1.data.experimental.shuffle_and_repeat tensorflow.compat.v1.size tensorflow.compat.v1.to_float tensorflow.compat.v1.random_uniform tensorflow.compat.v1.FixedLenSequenceFeature tensorflow.compat.v1.range tensorflow.compat.v1.greater_equal tensorflow.compat.v1.zeros tensorflow.compat.v1.matching_files tensorflow.compat.v1.ceil tensorflow.compat.v1.rank tensorflow.compat.v1.squeeze tensorflow.compat.v1.fill tensorflow.compat.v1.assert_greater_equal tensorflow.compat.v1.logging.info
apis: tensorflow.compat.v2.reduce_sum tensorflow.compat.v2.sign numpy.float64 numpy.round numpy.abs tensorflow.compat.v2.math.exp tensorflow.compat.v2.GradientTape numpy.arange numpy.log tensorflow.compat.v2.zeros_like tensorflow.compat.v2.enable_v2_behavior tensorflow.compat.v2.ones_like numpy.array numpy.concatenate tensorflow.compat.v2.constant tensorflow.compat.v2.keras.optimizers.SGD tensorflow.compat.v2.Variable tensorflow.compat.v2.abs
apis: tensorflow.compat.v2.nn.softmax_cross_entropy_with_logits tensorflow.compat.v2.reduce_sum tensorflow.compat.v2.transpose tensorflow.compat.v2.stop_gradient tensorflow.compat.v2.stack tensorflow.compat.v2.shape tensorflow.compat.v2.reduce_mean tensorflow.compat.v2.reshape tensorflow.compat.v2.square tensorflow.compat.v2.range tensorflow.compat.v2.split tensorflow.compat.v2.matmul