Skip to content

Instantly share code, notes, and snippets.

@raingo
raingo / ops.py
Last active March 30, 2021 16:50
multi dimensional softmax with tensorflow
import tensorflow as tf
"""
Multi dimensional softmax,
refer to https://github.com/tensorflow/tensorflow/issues/210
compute softmax along the dimension of target
the native softmax only supports batch_size x dimension
"""
def softmax(target, axis, name=None):