tensorflow_probability/python/distributions/vector_diffeomixture.py
def softmax(x, axis, name=None):
  """Equivalent to tf.nn.softmax but works around b/70297725."""
  with tf.name_scope(name or "softmax"):
    x = tf.convert_to_tensor(value=x, name="x")
    ndims = (
        tensorshape_util.rank(x.shape)
        if tensorshape_util.rank(x.shape) is not None else tf.rank(
            x, name="ndims"))
    axis = tf.convert_to_tensor(value=axis, dtype=tf.int32, name="axis")
    axis_ = tf.get_static_value(axis)
    if axis_ is not None:
      axis = np.int(ndims + axis_ if axis_ < 0 else axis_)
    else:
      axis = tf.where(axis < 0, ndims + axis, axis)
  return tf.nn.softmax(x, axis=axis)
def softmax(x, axis, name=None):
  """Equivalent to tf.nn.softmax but works around b/70297725."""
  with tf.name_scope(name or "softmax"):
    x = tf.convert_to_tensor(value=x, name="x")
    ndims = (
        tensorshape_util.rank(x.shape)
        if tensorshape_util.rank(x.shape) is not None else tf.rank(
            x, name="ndims"))
    axis = tf.convert_to_tensor(value=axis, dtype=tf.int32, name="axis")
    axis_ = tf.get_static_value(axis)
    if axis_ is not None:
      axis = np.int(ndims + axis_ if axis_ < 0 else axis_)
    else:
      axis = tf.where(axis < 0, ndims + axis, axis)
  return tf.nn.softmax(x, axis=axis)
def softmax ( x , axis , name = None ) : with tf . name_scope ( name or "softmax" ) : x = tf . convert_to_tensor ( value = x , name = "x" ) ndims = ( tensorshape_util . rank ( x . shape ) if tensorshape_util . rank ( x . shape ) is not None else tf . rank ( x , name = "ndims" ) ) axis = tf . convert_to_tensor ( value = axis , dtype = tf . int32 , name = "axis" ) axis_ = tf . get_static_value ( axis ) if axis_ is not None : axis = np . int ( ndims + axis_ if axis_ < 0 else axis_ ) else : axis = tf . where ( axis < 0 , ndims + axis , axis ) return tf . nn . softmax ( x , axis = axis )
Equivalent to tf.nn.softmax but works around b/70297725.
Equivalent to tf . nn . softmax but works around b / 70297725 .