Automatically compute batch size.

  Args:
    sequence_length: an integer
    mesh_shape: an input to mtf.convert_to_shape()
    layout_rules: an input to mtf.convert_to_layout_rules()
    tokens_per_split: an integer
  Returns:
    an integer
172514
def auto_batch_size(sequence_length,
                    mesh_shape,
                    layout_rules,
                    tokens_per_split=2048):
  
  num_splits = mtf.tensor_dim_to_mesh_dim_size(
      layout_rules, mesh_shape, mtf.Dimension("batch", 0))
  ret = max(1, tokens_per_split // sequence_length) * num_splits
  tf.logging.info(
      "AUTO_BATCH_SIZE tokens_per_split=%s num_splits=%s"
      " sequence_length=%s batch_size=%s"
      % (tokens_per_split, num_splits, sequence_length, ret))
  return ret
1341304
def original_query_sequence_length(self):
    
    if not self.is_aligned() or not self.entries.cigar:
      return self.query_sequence_length # take the naive approach
    # we are here with something aligned so take more intelligent cigar apporach
    return sum([x[0] for x in self.cigar_array if re.match('[HMIS=X]',x[1])])
532039
def chunks(seq, chunk_size):
    # type: (Sequence[T], int) -> Iterable[Sequence[T]]
    
    return (seq[i:i + chunk_size] for i in range(0, len(seq), chunk_size))
861167
def batch_fun(dataset, training, shapes, target_names, num_devices,
              batch_size_per_device=32, batch_size=None, eval_batch_size=32,
              bucket_length=32, buckets=None,
              batch_shuffle_size=128, max_eval_length=None):
  """"""
  del target_names
  # Batch size is batch_size_per_device * num_devices unless given directly.
  batch_size = batch_size or batch_size_per_device * num_devices
  # If bucketing is not specified, check if target shapes are variable.
  cur_batch_size = batch_size if training else eval_batch_size
  # Make cur_batch_size divisible by num_devices.
  cur_batch_size = max(cur_batch_size // num_devices, 1) * num_devices
  # Create heuristic buckets is none are specified.
  if buckets is None:
    variable_target_shapes = False
    target_shape = shapes[1]
    for dim in target_shape:
      if dim is None:
        variable_target_shapes = True
    tf.logging.info("Heuristically setting bucketing to %s based on shapes "
                    "of target tensors." % variable_target_shapes)
    if variable_target_shapes:
      bucket_boundaries = [bucket_length // 4, bucket_length // 2,
                           bucket_length, bucket_length * 2,
                           bucket_length * 4, bucket_length * 8,
                           bucket_length * 16]
      # We will pad to boundaries which pads to bucket_boundary - 1: add 1 here.
      bucket_boundaries = [b + 1 for b in bucket_boundaries]
      if not training:
        max_eval_length = max_eval_length or bucket_length * 32
        bucket_boundaries[-1] = max_eval_length
      bucket_batch_sizes = [cur_batch_size * 4, cur_batch_size * 2,
                            cur_batch_size, cur_batch_size // 2,
                            cur_batch_size // 4, cur_batch_size // 8,
                            cur_batch_size // 16, 1]
      if not training:
        bucket_batch_sizes[-2] = cur_batch_size // max_eval_length
      # Make batch sizes divisible by num_devices.
      bucket_batch_sizes = [max(b // num_devices, 1) * num_devices
                            for b in bucket_batch_sizes]
      buckets = (bucket_boundaries, bucket_batch_sizes)

  if buckets:
    tf.logging.info("Bucketing with buckets %s." % str(buckets))
    def example_length(_, target):
      return tf.shape(target)[0]
    boundaries, batch_sizes = buckets
    dataset = dataset.apply(tf.data.experimental.bucket_by_sequence_length(
        example_length, boundaries, batch_sizes,
        pad_to_bucket_boundary=True))
  else:
    dataset = dataset.padded_batch(cur_batch_size, shapes)
  if training:
    return dataset.shuffle(batch_shuffle_size)
  return dataset
1370818
def _get_example_length(example):
  """"""
  length = tf.maximum(tf.shape(example[0])[0], tf.shape(example[1])[0])
  return length
1092479
def fluent_shape(self) -> Sequence[int]:
        ''''''
        return tuple(self._shape.as_list()[1:] if self._batch else self._shape.as_list()[:])
792064
def length_from_embedding(emb):
  
  return tf.cast(tf.reduce_sum(mask_from_embedding(emb), [1, 2, 3]), tf.int32)
1143147
def generate_flat_data(self):
    
    # Construct the string statements.
    all_statements = []
    all_targets = []
    self.sequence_sizes_in = []
    self.sequence_sizes_out = []
    for _ in six.moves.range(self._batch_size):
      length, nest = self.curriculum_obj.fetch()
      seq_size_in = self._max_seq_length
      # Generate batch within max length.
      is_valid_sample = False
      tries_remaining = 10
      while not is_valid_sample:
        value, code = generate_code(length, nest, self._ops)
        tokens_in, seq_size_in = self.tokenize(
            code, self._max_seq_length, self._token_by_char)
        tokens_out, seq_size_out = self.tokenize(
            value, self._max_seq_length, self._token_by_char)
        is_valid_sample = self._max_seq_length >= seq_size_in
        if is_valid_sample:
          self.sequence_sizes_in.append(seq_size_in)
          self.sequence_sizes_out.append(seq_size_out)
        if tries_remaining == 0:
          raise ValueError("Could not generate a sample below the allowable "
                           "maximum, consider reducing either max_length or "
                           "max_nest.")
        else:
          tries_remaining -= 1
      all_statements += tokens_in
      all_targets += tokens_out
    # Store the flattened data.
    self.flat_data = np.array(all_statements, dtype=np.int64)
    self.num_tokens = self.flat_data.shape[0]
    self.flat_targets = np.array(all_targets, dtype=np.int64)
    self.num_tokens_target = self.flat_targets.shape[0]
    self.start_token = np.array(self.tokenize(
        [get_start_token()], 1)[0], dtype=np.int64)
    self.end_token = np.array(self.tokenize(
        [get_end_token()], 1)[0], dtype=np.int64)
1212565
def sizes(count, offset=0, max_chunk=500):
    
    if count is None:
        chunk = max_chunk
        while True:
            yield chunk, offset
            offset += chunk
    else:
        while count:
            chunk = min(count, max_chunk)
            count = max(0, count - max_chunk)
            yield chunk, offset
            offset += chunk
103670
def get_chunks(sequence, chunk_size):
    
    return [
        sequence[idx:idx + chunk_size]
        for idx in range(0, len(sequence), chunk_size)
    ]
1003816
def transform_to_length(nndata, length):
        

        if length is None:
            return nndata

        if length:
            for cn in range(length):
                if cn not in nndata.cn_weights:
                    nndata.cn_weights[cn] = 0
                    nndata.cn_nninfo[cn] = []

        return nndata
853461
def nt_yielder(self, graph, size):
        
        for grp in self.make_batch(size, graph):
            tmpg = Graph()
            # Add statements as list to tmp graph
            tmpg += grp
            yield (len(tmpg), tmpg.serialize(format='nt'))
639395
def _param_types_to_shape(self, param_types: Optional[str]) -> Sequence[int]:
        ''''''
        param_types = [] if param_types is None else param_types
        shape = tuple(self.object_table[ptype]['size'] for ptype in param_types)
        return shape
1050471
def _fluent_size(self, fluents, ordering) -> Sequence[Sequence[int]]:
        '''
        '''
        shapes = []
        for name in ordering:
            fluent = fluents[name]
            shape = self._param_types_to_shape(fluent.param_types)
            shapes.append(shape)
        return tuple(shapes)
560368
def action_size(self) -> Sequence[Shape]:
        ''''''
        return self._sizes(self._compiler.rddl.action_size)
1123104
def grouped_impl(wrap, size, sequence):
    
    iterator = iter(sequence)
    try:
        while True:
            batch = islice(iterator, size)
            yield list(chain((wrap(next(batch)),), batch))
    except StopIteration:
        return
728449
def generateSequences(n=2048, w=40, sequenceLength=5, sequenceCount=2,
                      sharedRange=None, seed=42):
  
  # Lots of room for noise sdrs
  patternAlphabetSize = 10*(sequenceLength * sequenceCount)
  patternMachine = PatternMachine(n, w, patternAlphabetSize, seed)
  sequenceMachine = SequenceMachine(patternMachine, seed)
  numbers = sequenceMachine.generateNumbers(sequenceCount, sequenceLength,
                                            sharedRange=sharedRange )
  generatedSequences = sequenceMachine.generateFromNumbers(numbers)

  return sequenceMachine, generatedSequences, numbers