Keras で特定の軸方向に softmax したい

Keras で特定の軸方向に softmax したいときとかあると思います。

f:id:cookie-box:20170815221557p:plain:w720

Keras のドキュメントによるとどの軸に沿って正規化するかを指定する axis というのを渡せるようです。
Activations - Keras Documentation
ただ、渡し方がよくわかりません。以下の箇所に axis=2 と書いてもエラーになります。

from keras.models import Sequential
from keras.layers import Dense, Reshape, Activation

if __name__ == '__main__':
  model = Sequential()
  model.add(Dense(3*4*5, activation='relu', input_shape=(10,)))
  model.add(Reshape((3, 4, 5)))
  model.add(Activation('softmax', axis=2)) # ERROR
  model.compile(loss='mean_squared_error', optimizer='rmsprop')
  print(model.summary())

しょうがないので softmax をするだけの自作レイヤーをつくります。これはエラーになりません。

from keras.engine.topology import Layer
from keras.models import Sequential
from keras.layers import Dense, Reshape
from keras.activations import softmax

# いまやりたい softmax をするためだけのレイヤー
class MySoftmax(Layer):
  def __init__(self, **kwargs):
    super(MySoftmax, self).__init__(**kwargs)
  def call(self, x):
    return(softmax(x, axis=2))

if __name__ == '__main__':
  model = Sequential()
  model.add(Dense(3*4*5, activation='relu', input_shape=(10,)))
  model.add(Reshape((3, 4, 5)))
  model.add(MySoftmax())
  model.compile(loss='mean_squared_error', optimizer='rmsprop')
  print(model.summary())

本当に希望の軸方向に正規化されているか確認してみます。何でもよいので適当な入力を流し込めばよいです。

  x, y = [], []
  # 1だらけの適当なデータ
  for i in range(20):
    x.append(numpy.ones((10)))
    y.append(numpy.ones((3,4,5)))
  x = numpy.array(x)
  y = numpy.array(y)
  
  model.fit(x, y)
  y_test = model.predict(x)
  print(y_test[0])
[[[ 0.21292862  0.24584062  0.2997719   0.15851443  0.25      ]
  [ 0.30612475  0.24584062  0.22358862  0.32254457  0.25      ]
  [ 0.16930516  0.24584062  0.22358862  0.25293246  0.25      ]
  [ 0.31164148  0.26247811  0.25305092  0.26600859  0.25      ]]

 [[ 0.18996513  0.24999966  0.18364088  0.25715023  0.3020851 ]
  [ 0.23858465  0.25000107  0.27690381  0.21845251  0.27901751]
  [ 0.25863758  0.24999966  0.35581443  0.30594474  0.19078693]
  [ 0.31281266  0.24999966  0.18364088  0.21845251  0.22811046]]

 [[ 0.21095483  0.36942184  0.22160934  0.2253582   0.24394487]
  [ 0.26182058  0.22892946  0.22160934  0.32392535  0.24394487]
  [ 0.31626979  0.159567    0.335172    0.2253582   0.24394487]
  [ 0.21095483  0.24208161  0.22160934  0.2253582   0.26816541]]]

出力の shape が (3, 4, 5) で、axis=2 としたので、5つの要素の和を取る方向がこれで正規化されて…いません。
和が1になったのは4つの要素の和を取る方向でした。
axis=1, 2, 3 についてどの軸の方向が正規化されるのか調べると以下でした。5つの要素の和を取る方向を正規化したかったら axis=3 とすべきようです。

f:id:cookie-box:20170815225335p:plain:w750

axis=0, -1 のときは以下でした。axis=0 は 4x5 の面内の和が1になりました。axis=-1 はむしろ当初やりたかった softmax が達成されていました。-1 はデフォルトのはずで、じゃあわざわざカスタムレイヤつくる必要なかったのでしょうか。
f:id:cookie-box:20170815230442p:plain:w500

これでやりたい softmax ができたので、今度はこの softmax に合わせた cross entropy を用意しなくては、と思っているのですがもしかしてそれもわざわざやらなくてよいのでしょうか。