code: def maybe_add_microbatch_axis(
      x: tf.Tensor,
      num_microbatches: Optional[type_aliases.BatchSize],
  ) -> tf.Tensor:
      """Adds the microbatch axis.
  
      Args:
        x: the input tensor.
        num_microbatches: If None, x is returned unchanged. Otherwise, must divide
          the batch size.
  
      Returns:
        The input tensor x, reshaped from [batch_size, ...] to
        [num_microbatches, batch_size / num_microbatches, ...].
      """
      if num_microbatches is None:
          return x
      with tf.control_dependencies(
          [
              tf.assert_equal(
                  tf.math.floormod(tf.shape(x)[0], num_microbatches), 0
              )
          ]
      ):
          return tf.reshape(
              x,
              tf.concat(
                  [[num_microbatches, -1], tf.shape(x)[1:]], axis=0
              ),
          )
  
label_desc: Use image node_id 1 for calling a global function with example usage: maybe_add_microbatch_axis(x, num_microbatches) and returns: x, tf