code: def _stochastically_use_non_top_expert(
      gate_logits, experts_dim, hparams
  ):
      """With a specified probability use the second place or lower experts."""
      # With the specified probability use the second place expert in place of the
      # top expert.
      tf.logging.info(
          "Using second place expert with prob: {}".format(
              hparams.moe_use_second_place_expert_prob
          )
      )
      _, top_expert_index = mtf.top_1(
          gate_logits, reduced_dim=experts_dim
      )
      top_expert_mask = mtf.one_hot(
          top_expert_index, experts_dim, dtype=gate_logits.dtype
      )
  
      # With probability moe_expert_use_second_place_expert_prob send the token to
      # the non-top expert.
      use_second_place_expert = mtf.cast(
          mtf.less(
              mtf.random_uniform(
                  gate_logits.mesh, gate_logits.shape[:-1]
              ),
              hparams.moe_use_second_place_expert_prob,
          ),
          gate_logits.dtype,
      )
      # Mask out the top logit.
      second_place_gate_logits = -1e9 * top_expert_mask + gate_logits
  
      # If a temperature is specified sample from the remaining N-1 experts.
      if hparams.moe_use_second_place_expert_prob_temp is not None:
          tf.logging.info(
              "Expert second place temp: {}".format(
                  hparams.moe_use_second_place_expert_prob_temp
              )
          )
          # What expert should be used.
          second_expert_index = mtf.sample_with_temperature(
              second_place_gate_logits,
              experts_dim,
              temperature=hparams.moe_use_second_place_expert_prob_temp,
          )
          second_expert_mask = mtf.one_hot(
              second_expert_index, experts_dim, dtype=gate_logits.dtype
          )
          # Set all logits to -inf that are not the sampled expert
          second_place_gate_logits += (1 - second_expert_mask) * -1e9
      gate_logits = (
          use_second_place_expert * second_place_gate_logits
          + (1 - use_second_place_expert) * gate_logits
      )
      return gate_logits
  
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 3 for calling a global function with example usage: _stochastically_use_non_top_expert(gate_logits, experts_dim, hparams) and returns: gate_logits