code: def adaptive_softmax_loss_fn(
      decoder: transformer.Unitransformer,
      context: transformer.Context,
      logits: mtf.Tensor,
      targets: mtf.Tensor,
      output_vocab_dim: mtf.Dimension,
  ) -> mtf.Tensor:
      """Custom loss to use when training with an adaptive softmax.
  
      Embedding and softmax weights must be shared in order for this function to
      work. Note that label smoothing and loss_on_targets_only is not supported and
      will be silently ignored.
  
      Args:
        decoder: a transformer.Unitransformer
        context: a transformer.Context
        logits: an mtf.Tensor, note that this will actually be the hidden state of
          the final decoder layer
        targets: an mtf.Tensor
        output_vocab_dim: an mtf.Dimension
  
      Returns:
        the loss
      """
      del output_vocab_dim
      hidden = logits
      vocab_embedding = context.shared_params["embedding"]
      return vocab_embedding.compute_loss(
          decoder, hidden=hidden, targets=targets, context=context
      )
  
imports: import math
  from typing import Dict, Sequence, Union
  import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import transformer
  from mesh_tensorflow.transformer import vocab_embeddings
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 1 for calling a global function with example usage: adaptive_softmax_loss_fn(decoder, context, logits, targets, output_vocab_dim) and returns: vocab_embedding