import functools
import tensorflow as tf

from .util import static_shape, static_rank, funcname_scope, \
                  dimension_indices, dtype_min, dtype_epsilon
from .control import unpack_to_array
from .mask import mask



@funcname_scope
def softmax(vector, sequence_length=None):
  assert static_rank(vector) == 2

  return tf.nn.softmax(vector) if sequence_length is None else \
         _dynamic_softmax(vector, sequence_length)


@funcname_scope
def _dynamic_softmax(vector, sequence_length):
  vector_length = tf.shape(vector)[1]
  mask_ = tf.cast(mask(sequence_length, vector_length), vector.dtype)
  vector_with_min = mask_ * vector + (1 - mask_) * dtype_min(vector.dtype)

  unnormal_dist = tf.exp(vector_with_min - _batch_max(vector_with_min)) * mask_
  return _batch_vector_scalar_div(
      unnormal_dist,
      tf.reduce_sum(unnormal_dist, [1]) + dtype_epsilon(unnormal_dist.dtype))


@funcname_scope
def _batch_vector_scalar_div(vector, scalar):
  assert static_rank(vector) == 2 and static_rank(scalar) == 1
  return vector / tf.expand_dims(scalar, 1)


@funcname_scope
def _batch_max(x):
  return tf.reduce_max(x, dimension_indices(x)[1:])

@@ -1,21 +1,4 @@
-import functools%0A
 impo
@@ -28,20 +28,16 @@
 om .
-util
  import 
 stat
@@ -36,21 +36,31 @@
 ort 
-static_shape,
+batch%0Afrom .util import
  sta
@@ -89,108 +89,32 @@
 pe, 
-%5C%0A                  dimension_indices, dtype_min, dtype_epsilon%0Afrom .control import unpack_to_array
+dtype_min, dtype_epsilon
 %0Afro
@@ -415,46 +415,8 @@
 h):%0A
-  vector_length = tf.shape(vector)%5B1%5D%0A
   ma
@@ -455,21 +455,27 @@
 th, 
-vector_length
+tf.shape(vector)%5B1%5D
 ), v
@@ -607,18 +607,42 @@
 _min
+%0A                        
  - 
-_
 batch
-_
+.
 max(
@@ -656,16 +656,32 @@
 with_min
+, keep_dims=True
 )) * mas
@@ -683,16 +683,17 @@
 * mask_%0A
+%0A
   return
@@ -697,389 +697,126 @@
 urn 
-_batch_vector_scalar_div(%0A      unnormal_dist,%0A      tf.reduce_sum(unnormal_dist, %5B1%5D) + dtype_epsilon(unnormal_dist.dtype))%0A%0A%0A@funcname_scope%0Adef _batch_vector_scalar_div(vector, scalar):%0A  assert static_rank(vector) == 2 and static_rank(scalar) == 1%0A  return vector / tf.expand_dims(scalar, 1)%0A%0A%0A@funcname_scope%0Adef _batch_max(x):%0A  return tf.reduce_max(x, dimension_indices(x)%5B1:%5D
+unnormal_dist / (batch.sum(unnormal_dist, keep_dims=True)%0A                          + dtype_epsilon(unnormal_dist.dtype)
 )%0A

Fix a bug in and refactor softmax module