parameters: (logits, compute_op, dim=-1, name=None)
return_statement: return output
docstring: Helper function for softmax and log_softmax.
  
    It reshapes and transposes the input logits into a 2-D Tensor and then invokes
    the tf.nn._softmax or tf.nn._log_softmax function. The output would be
    transposed and reshaped back.
  
    Args:
      logits: A non-empty `Tensor`. Must be one of the following types: `half`,
        `float32`, `float64`.
      compute_op: Either gen_nn_ops._softmax or gen_nn_ops._log_softmax
      dim: The dimension softmax would be performed on. The default is -1 which
        indicates the last dimension.
      name: A name for the operation (optional).
  
    Returns:
      A `Tensor`. Has the same type as `logits`. Same shape as `logits`.
    Raises:
      InvalidArgumentError: if `logits` is empty or `dim` is beyond the last
        dimension of `logits`.
docstring_summary: Helper function for softmax and log_softmax.
docstring_tokens: Helper function for softmax and log_softmax .
function: def _softmax(logits, compute_op, dim=-1, name=None):
    """Helper function for softmax and log_softmax.
  
    It reshapes and transposes the input logits into a 2-D Tensor and then invokes
    the tf.nn._softmax or tf.nn._log_softmax function. The output would be
    transposed and reshaped back.
  
    Args:
      logits: A non-empty `Tensor`. Must be one of the following types: `half`,
        `float32`, `float64`.
      compute_op: Either gen_nn_ops._softmax or gen_nn_ops._log_softmax
      dim: The dimension softmax would be performed on. The default is -1 which
        indicates the last dimension.
      name: A name for the operation (optional).
  
    Returns:
      A `Tensor`. Has the same type as `logits`. Same shape as `logits`.
    Raises:
      InvalidArgumentError: if `logits` is empty or `dim` is beyond the last
        dimension of `logits`.
    """
    def _swap_axis(logits, dim_index, last_index):
      """Swaps logits's dim_index and last_index."""
      return array_ops.transpose(logits, array_ops.concat(
          0, [math_ops.range(dim_index), [last_index],
              math_ops.range(dim_index + 1, last_index), [dim_index]]))
  
    logits = ops.convert_to_tensor(logits)
    if logits.get_shape().ndims is 2 and dim is -1:
      return compute_op(logits, name=name)
  
    # We need its original shape for shape inference.
    shape = logits.get_shape()
  
    # If dim is the last dimension, simply reshape the logits to a matrix and
    # apply the internal softmax.
    if dim is -1:
      input_shape = array_ops.shape(logits)
      logits = _flatten_outer_dims(logits)
      output = compute_op(logits, name=name)
      output = array_ops.reshape(output, input_shape)
      return output
  
    # If dim is not the last dimension, we have to do a reshape and transpose so
    # that we can still perform softmax on its last dimension.
  
    # Swap logits' dimension of dim and its last dimension.
    input_rank = array_ops.rank(logits)
    logits = _swap_axis(logits, dim, math_ops.sub(input_rank, 1))
    shape_after_swap = array_ops.shape(logits)
  
    # Reshape logits into a matrix.
    logits = _flatten_outer_dims(logits)
  
    # Do the actual softmax on its last dimension.
    output = compute_op(logits, name=name)
  
    # Transform back the output tensor.
    output = array_ops.reshape(output, shape_after_swap)
    output = _swap_axis(output, dim, math_ops.sub(input_rank, 1))
  
    # Make shape inference work since reshape and transpose may erase its static
    # shape.
    output.set_shape(shape)
  
    return output
function_tokens: def _softmax ( logits , compute_op , dim = - 1 , name = None ) : def _swap_axis ( logits , dim_index , last_index ) : """Swaps logits's dim_index and last_index.""" return array_ops . transpose ( logits , array_ops . concat ( 0 , [ math_ops . range ( dim_index ) , [ last_index ] , math_ops . range ( dim_index + 1 , last_index ) , [ dim_index ] ] ) ) logits = ops . convert_to_tensor ( logits ) if logits . get_shape ( ) . ndims is 2 and dim is - 1 : return compute_op ( logits , name = name ) # We need its original shape for shape inference. shape = logits . get_shape ( ) # If dim is the last dimension, simply reshape the logits to a matrix and # apply the internal softmax. if dim is - 1 : input_shape = array_ops . shape ( logits ) logits = _flatten_outer_dims ( logits ) output = compute_op ( logits , name = name ) output = array_ops . reshape ( output , input_shape ) return output # If dim is not the last dimension, we have to do a reshape and transpose so # that we can still perform softmax on its last dimension. # Swap logits' dimension of dim and its last dimension. input_rank = array_ops . rank ( logits ) logits = _swap_axis ( logits , dim , math_ops . sub ( input_rank , 1 ) ) shape_after_swap = array_ops . shape ( logits ) # Reshape logits into a matrix. logits = _flatten_outer_dims ( logits ) # Do the actual softmax on its last dimension. output = compute_op ( logits , name = name ) # Transform back the output tensor. output = array_ops . reshape ( output , shape_after_swap ) output = _swap_axis ( output , dim , math_ops . sub ( input_rank , 1 ) ) # Make shape inference work since reshape and transpose may erase its static # shape. output . set_shape ( shape ) return output