code: def _maybe_add_one_action(mask):
      """For time steps where the mask is all zeros, adds one action randomly."""
      batch_size = tf.shape(mask)[0]
      num_actions = tf.shape(mask)[1]
      extra_actions = tf.one_hot(
          tf.random.uniform(
              [batch_size], 0, num_actions, dtype=tf.int32
          ),
          depth=num_actions,
          dtype=tf.int32,
      )
      cond = tf.cast(tf.equal(tf.reduce_max(mask, axis=1), 1), tf.bool)
      return tf.compat.v1.where(cond, mask, extra_actions)
  
file_path: agents/tf_agents/bandits/environments/bernoulli_action_mask_tf_environment.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Callable
  import gin
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.environments import bandit_tf_environment
  from tf_agents.policies import utils
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import common
label_desc: Use image node_id 1 for calling a global function with example usage: _maybe_add_one_action(mask) and returns: tf