class_docstr: Compute standard deviation over features synchronously across replicas.
  
  Computes the mini batch standard deviation over the features of the previous
  layer at each batch
  by synchronizing the global batch statistics across all devices that are
  training the model. For specific details about batch normalization please
  refer to the `tf.keras.layers.MiniBatchStandardDeviationBase` layer docs.
  
  If this layer is used when using tf.distribute strategy to train models
  across devices/workers, there will be an allreduce call to aggregate batch
  statistics across all replicas at every training step. Without tf.distribute
  strategy, this layer behaves as a regular
  `tf.keras.layers.MiniBatchStandardDeviation`
  layer.
  
  Example usage:
  
  ```python
  strategy = tf.distribute.MirroredStrategy()
  
  with strategy.scope():
    model = tf.keras.Sequential()
    model.add(tf.keras.layers.Dense(16))
    model.add(SyncMiniBatchStandardDeviation())
  ```
class_docstr_tok: ["Compute","standard","deviation","over","features","synchronously","across","replicas",".","Computes","the","mini","batch","standard","deviation","over","the","features","of","the","previous","layer","at","each","batch","by","synchronizing","the","global","batch","statistics","across","all","devices","that","are","training","the","model",".","For","specific","details","about","batch","normalization","please","refer","to","the","`","tf.keras.layers.MiniBatchStandardDeviationBase","`","layer","docs",".","If","this","layer","is","used","when","using","tf.distribute","strategy","to","train","models","across","devices\/workers",",","there","will","be","an","allreduce","call","to","aggregate","batch","statistics","across","all","replicas","at","every","training","step",".","Without","tf.distribute","strategy",",","this","layer","behaves","as","a","regular","`","tf.keras.layers.MiniBatchStandardDeviation","`","layer",".","Example","usage",":","``","`","python","strategy","=","tf.distribute.MirroredStrategy","(",")","with","strategy.scope","(",")",":","model","=","tf.keras.Sequential","(",")","model.add","(","tf.keras.layers.Dense","(","16",")",")","model.add","(","SyncMiniBatchStandardDeviation","(",")",")","``","`"]
code: def _calculate_mean_feature_standard_deviation(
      self, inputs, batch_axis
  ):
      with tf.name_scope("mean_feature_standard_deviation"):
          # The dynamic range of fp16 is too limited to support the collection of
          # sufficient statistics. As a workaround we simply perform the operations
          # on 32-bit floats before converting the mean and variance back to fp16
          inputs_32 = (
              tf.cast(inputs, tf.float32)
              if inputs.dtype == tf.float16
              else inputs
          )
          replica_ctx = tf.distribute.get_replica_context()
          if replica_ctx:
              local_sum = tf.reduce_sum(
                  inputs_32, axis=batch_axis, keepdims=True
              )
              local_squared_sum = tf.reduce_sum(
                  tf.square(inputs_32), axis=batch_axis, keepdims=True
              )
              batch_size = tf.cast(
                  tf.shape(inputs_32)[batch_axis], tf.float32
              )
              # TODO(b/163099951): batch the all-reduces once we sort out the ordering
              # issue for NCCL. We don't have a mechanism to launch NCCL in the same
              # order in each replica nowadays, so we limit NCCL to batch all-reduces.
              y_sum = replica_ctx.all_reduce(
                  tf.distribute.ReduceOp.SUM, local_sum
              )
              y_squared_sum = replica_ctx.all_reduce(
                  tf.distribute.ReduceOp.SUM, local_squared_sum
              )
              global_batch_size = replica_ctx.all_reduce(
                  tf.distribute.ReduceOp.SUM, batch_size
              )
  
              mean = y_sum / global_batch_size
              y_squared_mean = y_squared_sum / global_batch_size
              # var = E(x^2) - E(x)^2
              variance = y_squared_mean - tf.square(mean)
              stddev = tf.sqrt(variance)
              mean_stddev = tf.reduce_mean(stddev, keepdims=True)
  
          else:
              mean_stddev = (
                  super()._calculate_mean_feature_standard_deviation(
                      inputs, batch_axis
                  )
              )
  
          return mean_stddev
  
file_path: graphics/tensorflow_graphics/projects/gan/minibatch_standarddeviation.py
imports: import tensorflow
label_desc: Use image node_id 1 for calling the SyncMiniBatchStandardDeviation obj's underlying member method code with example usage: obj._calculate_mean_feature_standard_deviation(inputs, batch_axis) and returns: mean_stddev