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)
Equivalent to tf.nn.softmax but works around b/70297725.