class_docstr: Keras' layer implementation for DynamicNormalization.
  
  Similar to Batch Normalization, DynamicNormalization normalizes the
  activations of the previous layer for each input.
  
  The normalization formula:
    ((|mean|^2 + p(mean)) / |x|^2 - 2 * g(mean)) / sqrt(variance)
  where
    p(mean) = prior_scale * mean + prior_offset
    g(mean) = mean_scale * mean + mean_offset.
  Here (prior_scale, prior_offset, mean_scale, mean_offset) are learnable
  parameters.
code: def call(self, inputs, training=None):
      training = self._get_training_value(training)
  
      inputs_dtype = inputs.dtype.base_dtype
      if inputs_dtype in (tf.float16, tf.bfloat16):
          # Do all math in float32 if given 16-bit inputs for numeric stability.
          # In particular, it's very easy for variance to overflow in float16 and
          # for safety we also choose to cast bfloat16 to float32.
          inputs = tf.cast(inputs, tf.float32)
  
      # Compute the axes along which to reduce the mean / variance
      input_shape = inputs.shape
      ndims = len(input_shape)
      reduction_axes = [i for i in range(ndims) if i not in self.axis]
  
      # Stops gradient update for the layers below in grow mode.
      # Intuitively when a new cluster is created, the gradients sent down to the
      # lower layers can disrupt the original weight, so it makes more sense to
      # freeze the other part when growing.
      inputs = tf.cond(
          tf.equal(self.mode, dm_ops.LOOKUP_WITH_GROW),
          lambda: tf.stop_gradient(inputs),
          lambda: inputs,
      )
  
      # Broadcasting only necessary for single-axis batch norm where the axis is
      # not the last dimension
      broadcast_shape = [1] * ndims
      broadcast_shape[self.axis[0]] = input_shape.dims[
          self.axis[0]
      ].value
  
      def _broadcast(v):
          if (
              v is not None
              and len(v.shape) != ndims
              and reduction_axes != list(range(ndims - 1))
          ):
              return tf.reshape(v, broadcast_shape)
          return v
  
      mean_scale, mean_offset = _broadcast(self.mean_scale), _broadcast(
          self.mean_offset
      )
      if not self.use_batch_normalization:
          prior_scale, prior_offset = _broadcast(
              self.prior_scale
          ), _broadcast(self.prior_offset)
  
      # Looks up mean and variances of from dynamic Gaussian memory.
      (
          self.mean,
          self.variance,
          self.distance,
          self.cluster_id,
      ) = dm_ops.dynamic_gaussian_memory_lookup(
          inputs,
          self.mode,
          self.dm_config,
          "dm_lookup",
          service_address=self.service_address,
      )
  
      if self.use_batch_normalization:
          outputs = tf.nn.batch_normalization(
              inputs,
              self.mean,
              self.variance,
              mean_offset,
              mean_scale,
              self.epsilon,
          )
      else:
          outputs = dynamic_normalization(
              inputs,
              self.mean,
              self.variance,
              prior_offset,
              mean_offset,
              prior_scale,
              mean_scale,
              self.epsilon,
          )
      # If some components of the shape got lost due to adjustments, fix that.
      outputs.set_shape(input_shape)
  
      return outputs
  
file_path: neural-structured-learning/research/carls/dynamic_normalization.py
imports: import typing
  from research.carls import dynamic_embedding_config_pb2
  from research.carls import dynamic_memory_ops
  import tensorflow
label_desc: Use image node_id 9 for calling the DynamicNormalization obj's underlying member method code with example usage: obj.call(inputs, training) and returns: outputs, v, tf