code: def _router_z_loss(
      logits, experts_dim, num_microbatches, importance=None
  ):
      """Loss that encourages router logits to remain small and improves stability.
  
      Args:
        logits: a tensor with shape [<batch_dims>, experts_dim]
        experts_dim: a Dimension (the number of experts)
        num_microbatches: number of microbatches
        importance: an optional tensor with shape [<batch_dims>, group_size_dim]
  
      Returns:
        z_loss: scalar loss only applied by non-padded tokens and normalized by
          num_microbatches.
      """
      log_z = mtf.reduce_logsumexp(logits, experts_dim)
      z_loss = mtf.square(log_z)
      if importance is not None:
          z_loss *= mtf.cast(
              mtf.equal(importance, 1.0), dtype=z_loss.dtype
          )
      denom = mtf.reduce_sum(
          mtf.cast(mtf.equal(importance, 1.0), dtype=z_loss.dtype)
      )
      z_loss = mtf.reduce_sum(z_loss) / (denom * num_microbatches)
      return z_loss
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import transformer
  import tensorflow.compat.v1
label_desc: Use image node_id 11 for calling a global function with example usage: _router_z_loss(logits, experts_dim, num_microbatches, importance) and returns: z_loss