Pad batch dim of features to nearest multiple of batch_multiple.
189268
def pad_batch(features, batch_multiple):
  """"""
  feature = list(features.items())[0][1]
  batch_size = tf.shape(feature)[0]
  mod = batch_size % batch_multiple
  has_mod = tf.cast(tf.cast(mod, tf.bool), tf.int32)
  batch_padding = batch_multiple * has_mod - mod

  padded_features = {}
  for k, feature in features.items():
    rank = len(feature.shape)
    paddings = [[0, 0] for _ in range(rank)]
    paddings[0][1] = batch_padding
    padded_feature = tf.pad(feature, paddings)
    padded_features[k] = padded_feature
  return padded_features
634109
def _fused_batch_norm_op(self, input_batch, mean, variance, use_batch_stats):
    """"""
    # Store the original shape of the mean and variance.
    mean_shape = mean.get_shape()
    variance_shape = variance.get_shape()
    # The fused batch norm expects the mean, variance, gamma and beta
    # tensors to have dimension 1, so we flatten them to remove the
    # extra dimensions. In addition, it expects the input_batch to have
    # dimension 4, so we reshape it accordingly.
    gamma_flatten = tf.reshape(self._gamma, shape=(self._num_channels,))
    beta_flatten = tf.reshape(self._beta, shape=(self._num_channels,))
    flatten_mean = tf.reshape(mean, shape=(self._num_channels,))
    flatten_variance = tf.reshape(variance, shape=(self._num_channels,))
    use_batch_stats = tf.convert_to_tensor(use_batch_stats)

    input_shape = input_batch.get_shape()
    output_shape = [-1] + input_shape.as_list()[1:]

    flat_image_size = np.prod(self._image_shape, dtype=np.int32)

    if len(self._data_format) == 4:
      fusable_data_format = self._data_format
      fusable_batch = input_batch
    elif self._channel_index == 1 and self._image_shape:
      fusable_data_format = "NCHW"
      fusable_batch = tf.reshape(
          input_batch,
          shape=(-1, self._num_channels, 1, flat_image_size))
    else:
      # The CPU implementation of FusedBatchNorm only supports NHWC tensor
      # format for now.
      fusable_data_format = "NHWC"
      fusable_batch = tf.reshape(
          input_batch,
          shape=(-1, 1, flat_image_size, self._num_channels))

    common_args = {
        "scale": gamma_flatten,
        "offset": beta_flatten,
        "epsilon": self._eps,
        "data_format": fusable_data_format,
        "name": "batch_norm"
    }

    def use_batch_stats_fused_batch_norm():
      return tf.nn.fused_batch_norm(
          fusable_batch,
          mean=None,
          variance=None,
          is_training=True,
          **common_args)

    def moving_average_fused_batch_norm():
      return tf.nn.fused_batch_norm(
          fusable_batch,
          mean=flatten_mean,
          variance=flatten_variance,
          is_training=False,
          **common_args)

    batch_norm_op, mean, variance = utils.smart_cond(
        use_batch_stats, use_batch_stats_fused_batch_norm,
        moving_average_fused_batch_norm)

    if len(self._data_format) != 4:
      batch_norm_op = tf.reshape(batch_norm_op, output_shape)
    mean = tf.reshape(mean, mean_shape)
    variance = tf.reshape(variance, variance_shape)
    return batch_norm_op, mean, variance
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
210041
def pad_conv3d_lrelu(self, activations, n_filters, kernel_size, strides,
                       scope):
    """"""
    padding = [[0, 0], [1, 1], [1, 1], [1, 1], [0, 0]]

    # tf.nn.conv3d accepts a list of 5 values for strides
    # with first and last value equal to 1
    if isinstance(strides, numbers.Integral):
      strides = [strides] * 3
    strides = [1] + strides + [1]

    # Filter_shape = [K, K, K, num_input, num_output]
    filter_shape = (
        [kernel_size]*3 + activations.shape[-1:].as_list() + [n_filters])

    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
      conv_filter = tf.get_variable(
          "conv_filter", shape=filter_shape,
          initializer=tf.truncated_normal_initializer(stddev=0.02))

      if self.hparams.use_spectral_norm:
        conv_filter, assign_op = common_layers.apply_spectral_norm(conv_filter)
        if self.is_training:
          tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, assign_op)

      padded = tf.pad(activations, padding)
      convolved = tf.nn.conv3d(
          padded, conv_filter, strides=strides, padding="VALID")
      rectified = tf.nn.leaky_relu(convolved, alpha=0.2)
    return rectified
98683
def cbow_fasttext_batch(centers, contexts, num_tokens, subword_lookup, dtype,
                        index_dtype):
    """"""
    _, contexts_row, contexts_col = contexts
    data, row, col = subword_lookup(contexts_row, contexts_col)
    centers = mx.nd.array(centers, dtype=index_dtype)
    contexts = mx.nd.sparse.csr_matrix(
        (data, (row, col)), dtype=dtype,
        shape=(len(centers), num_tokens))  # yapf: disable
    return centers, contexts
90531
def get_batch_coordinate(x, axis=0):
  """"""
  # Compute the batch coordinate before flattening all batches
  batch_coordinate = tf.expand_dims(
      common_attention.coordinate_tensor(tf.shape(x)[:-1], axis=axis), axis=-1)
  return batch_coordinate
90814
def batchify(data, batch_size):
    """"""
    nbatch = data.shape[0] // batch_size
    data = data[:nbatch * batch_size]
    data = data.reshape((batch_size, nbatch)).T
    return data
508353
def wider_bn(layer, start_dim, total_dim, n_add, weighted=True):
    '''
    '''
    n_dim = get_n_dim(layer)
    if not weighted:
        return get_batch_norm_class(n_dim)(layer.num_features + n_add)

    weights = layer.get_weights()

    new_weights = [
        add_noise(np.ones(n_add, dtype=np.float32), np.array([0, 1])),
        add_noise(np.zeros(n_add, dtype=np.float32), np.array([0, 1])),
        add_noise(np.zeros(n_add, dtype=np.float32), np.array([0, 1])),
        add_noise(np.ones(n_add, dtype=np.float32), np.array([0, 1])),
    ]

    student_w = tuple()
    for weight, new_weight in zip(weights, new_weights):
        temp_w = weight.copy()
        temp_w = np.concatenate(
            (temp_w[:start_dim], new_weight, temp_w[start_dim:total_dim])
        )
        student_w += (temp_w,)
    new_layer = get_batch_norm_class(n_dim)(layer.num_features + n_add)
    new_layer.set_weights(student_w)
    return new_layer
751603
def _split_batches(self, data, batch_size):
        """"""
        for i in range(0, len(data), batch_size):
            yield data[i : i + batch_size]
554637
def batch_fn(dataset, training, shapes, target_names,
             batch_size=32, eval_batch_size=32, bucket_batch_length=32,
             bucket_max_length=256, bucket_min_length=8,
             bucket_length_step=1.1, buckets=None):
  """"""
  del target_names
  # If bucketing is not specified, check if target shapes are variable.
  cur_batch_size = batch_size if training else eval_batch_size
  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:
      batch_size_per_token = cur_batch_size * bucket_batch_length
      scheme = data_reader.batching_scheme(batch_size_per_token,
                                           bucket_max_length,
                                           bucket_min_length,
                                           bucket_length_step,
                                           drop_long_sequences=training)
      buckets = (scheme["boundaries"], scheme["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))
  else:
    dataset = dataset.padded_batch(cur_batch_size, shapes)
  return dataset
1262873
def stub_batch(cls, size, **kwargs):
        
        return [cls.stub(**kwargs) for _ in range(size)]
194448
def _batch_norm_without_layers(self, input_layer, decay, use_scale,
                                   epsilon):
        """"""
        shape = input_layer.shape
        num_channels = shape[3] if self.data_format == "NHWC" else shape[1]
        beta = self.get_variable(
            "beta", [num_channels],
            tf.float32,
            tf.float32,
            initializer=tf.zeros_initializer())
        if use_scale:
            gamma = self.get_variable(
                "gamma", [num_channels],
                tf.float32,
                tf.float32,
                initializer=tf.ones_initializer())
        else:
            gamma = tf.constant(1.0, tf.float32, [num_channels])
        moving_mean = tf.get_variable(
            "moving_mean", [num_channels],
            tf.float32,
            initializer=tf.zeros_initializer(),
            trainable=False)
        moving_variance = tf.get_variable(
            "moving_variance", [num_channels],
            tf.float32,
            initializer=tf.ones_initializer(),
            trainable=False)
        if self.phase_train:
            bn, batch_mean, batch_variance = tf.nn.fused_batch_norm(
                input_layer,
                gamma,
                beta,
                epsilon=epsilon,
                data_format=self.data_format,
                is_training=True)
            mean_update = moving_averages.assign_moving_average(
                moving_mean, batch_mean, decay=decay, zero_debias=False)
            variance_update = moving_averages.assign_moving_average(
                moving_variance,
                batch_variance,
                decay=decay,
                zero_debias=False)
            tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, mean_update)
            tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, variance_update)
        else:
            bn, _, _ = tf.nn.fused_batch_norm(
                input_layer,
                gamma,
                beta,
                mean=moving_mean,
                variance=moving_variance,
                epsilon=epsilon,
                data_format=self.data_format,
                is_training=False)
        return bn
736935
def get_fc_dimensions(self, strides, kernel_sizes):
    """"""
    output_height, output_width, _ = self.hparams.problem.frame_shape
    output_steps = self.hparams.video_num_target_frames
    output_shape = np.array([output_steps, output_height, output_width])
    for curr_stride, kernel_size in zip(strides, kernel_sizes):
      output_shape = self.expected_output_shape(
          output_shape, np.array(curr_stride), 1, kernel_size)
    return np.prod(output_shape) * self.hparams.num_discriminator_filters * 8
85562
def with_dimensions(self, *dimensions):
    """"""
    self.dimensions = tuple(
        self._maybe_make_dimension(dim) for dim in dimensions)
    self._cached = None
    return self
1246104
def batch(input_iter, batch_size=32):
  """"""
  input_iter = iter(input_iter)
  next_ = list(itertools.islice(input_iter, batch_size))
  while next_:
    yield next_
    next_ = list(itertools.islice(input_iter, batch_size))
1094509
def _pooling_output_shape(input_shape, pool_size=(2, 2),
                          strides=None, padding='VALID'):
  """"""
  dims = (1,) + pool_size + (1,)  # NHWC
  spatial_strides = strides or (1,) * len(pool_size)
  strides = (1,) + spatial_strides + (1,)
  pads = padtype_to_pads(input_shape, dims, strides, padding)
  operand_padded = onp.add(input_shape, onp.add(*zip(*pads)))
  t = onp.floor_divide(onp.subtract(operand_padded, dims), strides) + 1
  return tuple(t)
436515
def _batch_norm(name, x):
  """"""
  with tf.name_scope(name):
    return tf.contrib.layers.batch_norm(
        inputs=x,
        decay=.9,
        center=True,
        scale=True,
        activation_fn=None,
        updates_collections=None,
        is_training=False)