Python tensorflow.nn() Examples

def embedding_lookup(embedding_matrix, indices, ids, weights, size):
  """Performs a weighted embedding lookup.

  Args:
    embedding_matrix: float Tensor from which to do the lookup.
    indices: int Tensor for the output rows of the looked up vectors.
    ids: int Tensor vectors to look up in the embedding_matrix.
    weights: float Tensor weights to apply to the looked up vectors.
    size: int number of output rows. Needed since some output rows may be
        empty.

  Returns:
    Weighted embedding vectors.
  """
  embeddings = tf.nn.embedding_lookup([embedding_matrix], ids)
  # TODO(googleuser): allow skipping weights.
  broadcast_weights_shape = tf.concat([tf.shape(weights), [1]], 0)
  embeddings *= tf.reshape(weights, broadcast_weights_shape)
  embeddings = tf.unsorted_segment_sum(embeddings, indices, size)
  return embeddings 

def _batch_conv_block(self, x, scope, x_dim, y_dim, in_channels, out_channels, conv_size = (3,3), pooling=False, pool_size=(2,2)):
        with tf.variable_scope(scope):
            
            x = tf.reshape(x, [self._batch_size, x_dim, y_dim, in_channels])
            h1 = tf.contrib.layers.conv2d(x, out_channels,
                                          conv_size,
                                          activation_fn=None,
                                          weights_regularizer=None)
            h2 = tf.contrib.layers.batch_norm(h1,
                                              center=True, scale=True,
                                              is_training=self._phase)

            # return tf.nn.relu(h2, 'relu')
            h3 = self._activation_fn(h2,name='activation_fn')

            #h3 = tf.layers.flatten(h3)

            #return h3
            if pooling:
                h3 = tf.contrib.layers.max_pool2d(h3, pool_size)

            #h4 = tf.layers.flatten(h3)

            return h3 

def multilayer_perceptron(final_output, weights, biases):
  """MLP over output with attention over enc outputs
  Args:
     final_output: [batch_size x 2*size]
  Returns:
     logit:  [batch_size x target_label_size]
  """
  
  # Layer 1
  layer_1 = tf.add(tf.matmul(final_output, weights["h1"]), biases["b1"])
  layer_1 = tf.nn.relu(layer_1)

  # Layer 2
  layer_2 = tf.add(tf.matmul(layer_1, weights["h2"]), biases["b2"])
  layer_2 = tf.nn.relu(layer_2)

  # output layer
  layer_out = tf.add(tf.matmul(layer_2, weights["out"]), biases["out"])
  
  return layer_out 

def simple_rnn(rnn_input, initial_state=None):
  """Implements Simple RNN
  Args:
  rnn_input: List of tensors of sizes [-1, sentembed_size]
  Returns:
  encoder_outputs, encoder_state
  """     
  # Setup cell
  cell_enc = get_lstm_cell()
  
  # Setup RNNs
  dtype = tf.float16 if FLAGS.use_fp16 else tf.float32
  rnn_outputs, rnn_state = tf.nn.rnn(cell_enc, rnn_input, dtype=dtype, initial_state=initial_state)
  # print(rnn_outputs)
  # print(rnn_state)
  
  return rnn_outputs, rnn_state 

def _loop_function(current_inp, ext_logits, gold_logits):
  """ Update current input wrt previous logits
  Args:
  current_inp: [batch_size x sentence_embedding_size]
  ext_logits: [batch_size x target_label_size] [1, 0]
  gold_logits: [batch_size x target_label_size]
  Returns:
  updated_inp: [batch_size x sentence_embedding_size]
  """

  prev_logits = gold_logits
  if not FLAGS.authorise_gold_label:
    prev_logits = ext_logits
    prev_logits = tf.nn.softmax(prev_logits) # [batch_size x target_label_size]

  prev_logits = tf.split(1, FLAGS.target_label_size, prev_logits) # [[batch_size], [batch_size], ...]
  prev_weight_one = prev_logits[0]
    
  updated_inp = tf.mul(current_inp, prev_weight_one)
  # print(updated_inp)

  return updated_inp


### SoftMax and Predictions 

def variable_summaries(var, name, collection_key):
    """Attach a lot of summaries to a Tensor (for TensorBoard visualization).

    Args:
        - var: Tensor for variable from which we want to log.
        - name: Variable name.
        - collection_key: Collection to save the summary to, can be any key of
          `VAR_LOG_LEVELS`.
    """
    if collection_key not in VAR_LOG_LEVELS.keys():
        raise ValueError('"{}" not in `VAR_LOG_LEVELS`'.format(collection_key))
    collections = VAR_LOG_LEVELS[collection_key]

    with tf.name_scope(name):
        mean = tf.reduce_mean(var)
        tf.summary.scalar('mean', mean, collections)
        num_params = tf.reduce_prod(tf.shape(var))
        tf.summary.scalar('num_params', num_params, collections)
        with tf.name_scope('stddev'):
            stddev = tf.sqrt(tf.reduce_mean(tf.square(var - mean)))
        tf.summary.scalar('stddev', stddev, collections)
        tf.summary.scalar('max', tf.reduce_max(var), collections)
        tf.summary.scalar('min', tf.reduce_min(var), collections)
        tf.summary.histogram('histogram', var, collections)
        tf.summary.scalar('sparsity', tf.nn.zero_fraction(var), collections) 

def get_rnn_cell_trainable_variables(cell):
    """Returns the list of trainable variables of an RNN cell.

    Args:
        cell: an instance of :tf_main:`RNNCell <nn/rnn_cell/RNNCell>`.

    Returns:
        list: trainable variables of the cell.
    """
    cell_ = cell
    while True:
        try:
            return cell_.trainable_variables
        except AttributeError:
        # Cell wrappers (e.g., `DropoutWrapper`) cannot directly access to
        # `trainable_variables` as they don't initialize superclass
        # (tf==v1.3). So try to access through the cell in the wrapper.
            cell_ = cell._cell  # pylint: disable=protected-access 

def conv_1d(inputs, weights, biases,
            stride=1, padding='SAME',
            activation='relu', norm=None,
            dropout=False, dropout_rate=None,
            is_training=True):
    hidden = tf.nn.conv1d(tf.cast(inputs, tf.float32), weights, stride=stride,
                          padding=padding) + biases

    if norm is not None:
        if norm == 'batch':
            hidden = tf.layers.batch_normalization(
                hidden,
                training=is_training
            )
        elif norm == 'layer':
            hidden = tf.contrib.layers.layer_norm(hidden)

    if activation:
        hidden = getattr(tf.nn, activation)(hidden)

    if dropout and dropout_rate is not None:
        hidden = tf.layers.dropout(hidden, rate=dropout_rate,
                                   training=is_training)

    return hidden 

def conv_2d(inputs, weights, biases,
            stride=1, padding='SAME',
            activation='relu', norm=None,
            dropout=False, dropout_rate=None,
            is_training=True):
    hidden = tf.nn.conv2d(inputs, weights, strides=[1, stride, stride, 1],
                          padding=padding) + biases

    if norm is not None:
        if norm == 'batch':
            hidden = tf.layers.batch_normalization(
                hidden,
                training=is_training
            )
        elif norm == 'layer':
            hidden = tf.contrib.layers.layer_norm(hidden)

    if activation:
        hidden = getattr(tf.nn, activation)(hidden)

    if dropout and dropout_rate is not None:
        hidden = tf.layers.dropout(hidden, rate=dropout_rate,
                                   training=is_training)

    return hidden 

def build_output(self, inputs, inferences):
    scores = tf.nn.softmax(inferences, name='scores')
    tf.add_to_collection('outputs', scores)

    with tf.name_scope('labels'):
      label_indices = tf.arg_max(inferences, 1, name='arg_max')
      labels = self.classification.output_labels(label_indices)
      tf.add_to_collection('outputs', labels)

    keys = self.classification.keys(inputs)
    if keys:
      # Key feature, if it exists, is a passthrough to the output.
      # The use of identity is to name the tensor and correspondingly the output field.
      keys = tf.identity(keys, name='key')
      tf.add_to_collection('outputs', keys)

    return {
      'label': labels,
      'score': scores
    } 

def embedding_lookup(embedding_matrix, indices, ids, weights, size):
  """Performs a weighted embedding lookup.

  Args:
    embedding_matrix: float Tensor from which to do the lookup.
    indices: int Tensor for the output rows of the looked up vectors.
    ids: int Tensor vectors to look up in the embedding_matrix.
    weights: float Tensor weights to apply to the looked up vectors.
    size: int number of output rows. Needed since some output rows may be
        empty.

  Returns:
    Weighted embedding vectors.
  """
  embeddings = tf.nn.embedding_lookup([embedding_matrix], ids)
  # TODO(googleuser): allow skipping weights.
  broadcast_weights_shape = tf.concat([tf.shape(weights), [1]], 0)
  embeddings *= tf.reshape(weights, broadcast_weights_shape)
  embeddings = tf.unsorted_segment_sum(embeddings, indices, size)
  return embeddings 

def efficientnet(width_coefficient=None,
                 depth_coefficient=None,
                 dropout_rate=0.2,
                 survival_prob=0.8):
  """Creates a efficientnet model."""
  global_params = efficientnet_model.GlobalParams(
      blocks_args=_DEFAULT_BLOCKS_ARGS,
      batch_norm_momentum=0.99,
      batch_norm_epsilon=1e-3,
      dropout_rate=dropout_rate,
      survival_prob=survival_prob,
      data_format='channels_last',
      num_classes=1000,
      width_coefficient=width_coefficient,
      depth_coefficient=depth_coefficient,
      depth_divisor=8,
      min_depth=None,
      relu_fn=tf.nn.swish,
      # The default is TPU-specific batch norm.
      # The alternative is tf.layers.BatchNormalization.
      batch_norm=utils.BatchNormalization,  # TPU-specific requirement.
      use_se=True,
      clip_projection_output=False)
  return global_params 

def build(self, img):
        """Constructs the generative network's layers. Normally called after initialization.
        
        Args:
            img: 4D tensor representation of image batch
        """

        self.padded = self._pad(img, 40)

        self.conv1 = self._conv_block(self.padded, maps_shape=[9, 9, 3, 32], stride=1, name='conv1')
        self.conv2 = self._conv_block(self.conv1, maps_shape=[2, 2, 32, 64], stride=2, name='conv2')
        self.conv3 = self._conv_block(self.conv2, maps_shape=[2, 2, 64, 128], stride=2, name='conv3')

        self.resid1 = self._residual_block(self.conv3, maps_shape=[3, 3, 128, 128], stride=1, name='resid1')
        self.resid2 = self._residual_block(self.resid1, maps_shape=[3, 3, 128, 128], stride=1, name='resid2')
        self.resid3 = self._residual_block(self.resid2, maps_shape=[3, 3, 128, 128], stride=1, name='resid3')
        self.resid4 = self._residual_block(self.resid3, maps_shape=[3, 3, 128, 128], stride=1, name='resid4')
        self.resid5 = self._residual_block(self.resid4, maps_shape=[3, 3, 128, 128], stride=1, name='resid5')

        self.conv4 = self._upsample_block(self.resid5, maps_shape=[2, 2, 64, 128], stride=2, name='conv4')
        self.conv5 = self._upsample_block(self.conv4, maps_shape=[2, 2, 32, 64], stride=2, name='conv5')
        self.conv6 = self._conv_block(self.conv5, maps_shape=[9, 9, 32, 3], stride=1, name='conv6', activation=None)

        self.output = tf.nn.sigmoid(self.conv6) 

def _instance_normalize(inputs):
        """Instance normalize inputs to reduce covariate shift and reduce dependency on input contrast to improve results.
        
        Args:
            inputs: 4D tensor representing image layer encodings
            
        Returns:
            maps: 4D tensor of batch normalized inputs
        """

        with tf.variable_scope('instance_normalization'):
            batch, height, width, channels = [_.value for _ in inputs.get_shape()]
            mu, sigma_sq = tf.nn.moments(inputs, [1, 2], keep_dims=True)

            shift = tf.Variable(tf.constant(.1, shape=[channels]))
            scale = tf.Variable(tf.ones([channels]))
            normalized = (inputs - mu) / (sigma_sq + EPSILON) ** .5
            maps = scale * normalized + shift
            return maps 

def embedding_lookup(embedding_matrix, indices, ids, weights, size):
  """Performs a weighted embedding lookup.

  Args:
    embedding_matrix: float Tensor from which to do the lookup.
    indices: int Tensor for the output rows of the looked up vectors.
    ids: int Tensor vectors to look up in the embedding_matrix.
    weights: float Tensor weights to apply to the looked up vectors.
    size: int number of output rows. Needed since some output rows may be
        empty.

  Returns:
    Weighted embedding vectors.
  """
  embeddings = tf.nn.embedding_lookup([embedding_matrix], ids)
  # TODO(googleuser): allow skipping weights.
  broadcast_weights_shape = tf.concat([tf.shape(weights), [1]], 0)
  embeddings *= tf.reshape(weights, broadcast_weights_shape)
  embeddings = tf.unsorted_segment_sum(embeddings, indices, size)
  return embeddings 

def log_prob(self, param_batch, sample_vecs):
        if hasattr(tf.nn, 'softmax_cross_entropy_with_logits_v2'):
            loss_func = tf.nn.softmax_cross_entropy_with_logits_v2
        else:
            loss_func = tf.nn.softmax_cross_entropy_with_logits
        return tf.negative(loss_func(labels=sample_vecs, logits=param_batch)) 

def _batch_relu_layer(self, x, size, scope):
        with tf.variable_scope(scope):
            h1 = tf.contrib.layers.fully_connected(x, size,
                                                   activation_fn=None,
                                                   weights_regularizer=None,
                                                   scope='dense')
            h2 = tf.contrib.layers.batch_norm(h1,
                                              center=True, scale=True,
                                              is_training=self._phase,
                                              scope='bn')
            # return tf.nn.relu(h2, 'relu')
            return self._activation_fn(h2, name='activation_fn') 

def _batch_relu_layer(self, x, scope, in_size, out_size):
        with tf.variable_scope(scope):
            x = tf.reshape(x, [self._batch_size, in_size])
            h1 = tf.contrib.layers.fully_connected(x, out_size,
                                                   activation_fn=None,
                                                   weights_regularizer=None)
            h2 = tf.contrib.layers.batch_norm(h1,
                                              center=True, scale=True,
                                              is_training=self._phase,
                                              scope='bn')
            # return tf.nn.relu(h2, 'relu')
            return self._activation_fn(h2,name='activation_fn') 

def create_estimator(params=None):
  """Returns the glyphs DNNClassifier estimator.

  Args:
    params: Optional hyperparameters, defaulting to command-line values.

  Returns:
    A DNNClassifier instance.
  """
  params = params or get_flag_params()
  if not params['layer_dims'] and params['activation_fn'] != 'sigmoid':
    tf.logging.warning(
        'activation_fn should be sigmoid for logistic regression. Got: %s',
        params['activation_fn'])

  activation_fn = getattr(tf.nn, params['activation_fn'])
  estimator = tf.estimator.DNNClassifier(
      params['layer_dims'],
      feature_columns=[glyph_patches.create_patch_feature_column()],
      weight_column=glyph_patches.WEIGHT_COLUMN_NAME,
      n_classes=len(musicscore_pb2.Glyph.Type.keys()),
      optimizer=tf.train.FtrlOptimizer(
          learning_rate=params['learning_rate'],
          l1_regularization_strength=params['l1_regularization_strength'],
          l2_regularization_strength=params['l2_regularization_strength'],
      ),
      activation_fn=activation_fn,
      dropout=FLAGS.dropout,
      model_dir=glyph_patches.FLAGS.model_dir,
  )
  return hyperparameters.estimator_with_saved_params(estimator, params) 

def _degree_matrix(self, vol_shape):
        # get shape stats
        ndims = len(vol_shape)
        sz = [*vol_shape, ndims]

        # prepare conv kernel
        conv_fn = getattr(tf.nn, 'conv%dd' % ndims)

        # prepare tf filter
        z = K.ones([1] + sz)
        filt_tf = tf.convert_to_tensor(self._adj_filt(ndims), dtype=tf.float32)
        strides = [1] * (ndims + 2)
        return conv_fn(z, filt_tf, strides, "SAME") 

def module_names():
  return [
      "tf",
      "tf.errors",
      "tf.image",
      "tf.nn",
      "tf.nn.rnn_cell",
      "tf.train",
      "tf.python_io",
      "tf.summary",
      "tf.test",
      "tf.contrib.bayesflow.entropy",
      "tf.contrib.bayesflow.monte_carlo",
      "tf.contrib.bayesflow.stochastic_graph",
      "tf.contrib.bayesflow.stochastic_tensor",
      "tf.contrib.bayesflow.variational_inference",
      "tf.contrib.copy_graph",
      "tf.contrib.crf",
      "tf.contrib.distributions",
      "tf.contrib.distributions.bijector",
      "tf.contrib.ffmpeg",
      "tf.contrib.framework",
      "tf.contrib.graph_editor",
      "tf.contrib.integrate",
      "tf.contrib.layers",
      "tf.contrib.learn",
      "tf.contrib.learn.monitors",
      "tf.contrib.losses",
      "tf.contrib.rnn",
      "tf.contrib.metrics",
      "tf.contrib.training",
      "tf.contrib.util",
  ] 

def get_lstm_cell():
  """Define LSTM Cell
  """
  single_cell = tf.nn.rnn_cell.BasicLSTMCell(FLAGS.size) if (FLAGS.lstm_cell == "lstm") else tf.nn.rnn_cell.GRUCell(FLAGS.size)
  cell = single_cell
  if FLAGS.num_layers > 1:
    cell = tf.nn.rnn_cell.MultiRNNCell([single_cell] * FLAGS.num_layers)
  return cell

### Reshaping 

def softmax(x):
  input_shape = tf.shape(x)
  input_rank = tf.shape(input_shape)[0]
  input_size = tf.gather(input_shape, input_rank-1)
  output_shape = input_shape
  
  x = tf.reshape(x, [-1, input_size])

  y = tf.nn.softmax(x)
  y = tf.reshape(y, output_shape)
  
  return y 

def convLayer(x, filter_size=5, filter_depth=64, pool_size=2):
  x_depth = x.get_shape()[-1].value
  W = weight_variable([filter_size, filter_size, x_depth, filter_depth])
  conv = tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')

  b = bias_variable([filter_depth])
  relu = tf.nn.relu(conv + b)

  pool = tf.nn.max_pool(relu,
                        ksize=[1,pool_size,pool_size,1],
                        strides=[1,pool_size,pool_size,1],
                        padding = 'SAME')

  return pool 

def max_pool_2x2(x):
    '''
    Genarates a max-pool TensorFlow Op. This Op "strides" a window across the input x. In each window, the maximum value
    is selected and chosen to represent that region in the output Tensor. Hence the size/dimensionality of the problem
    is reduced.
    :param x: A Tensor with dimensions [batch_size, height, width, 3]
    :return: A TensorFlow Op that max-pools the input Tensor, x.
    '''
    return tf.nn.max_pool(x, ksize=[1, 2, 2, 1],
                          strides=[1, 2, 2, 1], padding='SAME') 

def conv2d(x, W):
    '''
    Generates a conv2d TensorFlow Op. This Op flattens the weight matrix (filter) down to 2D, then "strides" across the
    input Tensor x, selecting windows/patches. For each little_patch, the Op performs a right multiply:
            W . little_patch
    and stores the result in the output layer of feature maps.
    :param x: a minibatch of images with dimensions [batch_size, height, width, 3]
    :param W: a "filter" with dimensions [window_height, window_width, input_channels, output_channels]
    e.g. for the first conv layer:
          input_channels = 3 (RGB)
          output_channels = number_of_desired_feature_maps
    :return: A TensorFlow Op that convolves the input x with the filter W.
    '''
    return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME') 

def __call__(self, x):
    if self.nl == 'leaky_relu':
      return leaky_relu(x, self.alpha)
    elif self.nl == 'leaky_softplus':
      return leaky_softplus(x, self.alpha)
    else:
      return getattr(tf.nn, self.nl)(x) 

def attention(self, last_layer, attention_tensor):
    """Compute the attention term for the network unit."""
    h_tensor = attention_tensor

    # Compute the attentions.
    # Using feed-forward net to map the two inputs into the same dimension
    focus_tensor = tf.nn.tanh(
        tf.matmul(
            h_tensor,
            self._component.get_variable('attention_weights_pm_0'),
            name='h_x_pm') + self._component.get_variable('attention_bias_0'))

    context_tensor = tf.nn.tanh(
        tf.matmul(
            last_layer,
            self._component.get_variable('attention_weights_hm_0'),
            name='l_x_hm') + self._component.get_variable('attention_bias_1'))
    # The tf.multiply in the following expression broadcasts along the 0 dim:
    z_vec = tf.reduce_sum(tf.multiply(focus_tensor, context_tensor), 1)
    p_vec = tf.nn.softmax(tf.reshape(z_vec, [1, -1]))
    # The tf.multiply in the following expression broadcasts along the 1 dim:
    r_vec = tf.expand_dims(
        tf.reduce_sum(
            tf.multiply(
                h_tensor, tf.reshape(p_vec, [-1, 1]), name='time_together2'),
            0), 0)
    return tf.matmul(
        r_vec,
        self._component.get_variable('attention_weights_pu'),
        name='time_together3') 

def cross_entropy_loss(logits, labels, weights):
    """Estimate cost of predictions
    Add summary for "cost" and "cost/avg".
    Args:
    logits: Logits from inference(). [FLAGS.batch_size, FLAGS.max_doc_length, FLAGS.target_label_size]
    labels: Sentence extraction gold levels [FLAGS.batch_size, FLAGS.max_doc_length, FLAGS.target_label_size]
    weights: Weights to avoid padded part [FLAGS.batch_size, FLAGS.max_doc_length]
    Returns:
    Cross-entropy Cost
    """
    with tf.variable_scope('CrossEntropyLoss') as scope:
        # Reshape logits and labels to match the requirement of softmax_cross_entropy_with_logits
        logits = tf.reshape(logits, [-1, FLAGS.target_label_size]) # [FLAGS.batch_size*FLAGS.max_doc_length, FLAGS.target_label_size]
        labels = tf.reshape(labels, [-1, FLAGS.target_label_size]) # [FLAGS.batch_size*FLAGS.max_doc_length, FLAGS.target_label_size]
        cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits, labels) # [FLAGS.batch_size*FLAGS.max_doc_length]
        cross_entropy = tf.reshape(cross_entropy, [-1, FLAGS.max_doc_length])  # [FLAGS.batch_size, FLAGS.max_doc_length]
        if FLAGS.weighted_loss:
            cross_entropy = tf.mul(cross_entropy, weights)
      
        # Cross entroy / document
        cross_entropy = tf.reduce_sum(cross_entropy, reduction_indices=1) # [FLAGS.batch_size]
        cross_entropy_mean = tf.reduce_mean(cross_entropy, name='crossentropy')

        # ## Cross entroy / sentence
        # cross_entropy_sum = tf.reduce_sum(cross_entropy)
        # valid_sentences = tf.reduce_sum(weights)
        # cross_entropy_mean = cross_entropy_sum / valid_sentences
        
        # cross_entropy = -tf.reduce_sum(labels * tf.log(logits), reduction_indices=1)
        # cross_entropy_mean = tf.reduce_mean(cross_entropy, name='crossentropy')
    
        tf.add_to_collection('cross_entropy_loss', cross_entropy_mean)
        # # # The total loss is defined as the cross entropy loss plus all of
        # # # the weight decay terms (L2 loss).
        # # return tf.add_n(tf.get_collection('losses'), name='total_loss')     
    return cross_entropy_mean

### Training functions 

def normalize(self, inputs):
    """Apply normalization to input.

    The shape must match the declared shape in the constructor.
    [This is copied from tf.contrib.rnn.LayerNormBasicLSTMCell.]

    Args:
      inputs: Input tensor

    Returns:
      Normalized version of input tensor.

    Raises:
      ValueError: if inputs has undefined rank.
    """
    inputs_shape = inputs.get_shape()
    inputs_rank = inputs_shape.ndims
    if inputs_rank is None:
      raise ValueError('Inputs %s has undefined rank.' % inputs.name)
    axis = range(1, inputs_rank)

    beta = self._component.get_variable('beta_%s' % self._name)
    gamma = self._component.get_variable('gamma_%s' % self._name)

    with tf.variable_scope('layer_norm_%s' % self._name):
      # Calculate the moments on the last axis (layer activations).
      mean, variance = nn.moments(inputs, axis, keep_dims=True)

      # Compute layer normalization using the batch_normalization function.
      variance_epsilon = 1E-12
      outputs = nn.batch_normalization(inputs, mean, variance, beta, gamma,
                                       variance_epsilon)
      outputs.set_shape(inputs_shape)
      return outputs