Skip to content

Commit e2d6511

Browse files
authored
[Bugfix][Frontend][Keras]Fix a corner case bug in softmax converter of keras frontend (#15337)
* Fix softmax converter about keras * add new test cases to capture the bug * Update keras.py
1 parent c4f10cd commit e2d6511

File tree

2 files changed

+11
-2
lines changed

2 files changed

+11
-2
lines changed

python/tvm/relay/frontend/keras.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -131,11 +131,13 @@ def _convert_advanced_activation(inexpr, keras_layer, etab, data_layout, input_s
131131

132132
if act_type == "Softmax":
133133
axis = keras_layer.axis
134-
dims = len(input_shape)
134+
dims = len(input_shape) if input_shape else 0
135135
if isinstance(axis, list):
136136
raise tvm.error.OpAttributeUnImplemented(f"Softmax with axes {axis} is not supported.")
137137
if data_layout == "NCHW":
138-
if axis == -1:
138+
if dims == 0:
139+
axis = 0
140+
elif axis == -1:
139141
axis = 1
140142
else:
141143
axis = axis + 1 if axis < dims - 1 else 1

tests/python/frontend/keras/test_forward.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,13 @@ def test_forward_activations(self, keras_mod):
229229
keras_model = keras_mod.models.Model(data, x)
230230
verify_keras_frontend(keras_model)
231231
verify_keras_frontend(keras_model, need_transpose=False, layout="NHWC")
232+
# Test the input dimension = 1
233+
data = keras_mod.layers.Input(shape=(11,))
234+
act_func = keras_mod.layers.Softmax()
235+
x = act_func(data)
236+
keras_model = keras_mod.models.Model(data, x)
237+
verify_keras_frontend(keras_model)
238+
verify_keras_frontend(keras_model, need_transpose=False, layout="NHWC")
232239

233240
def test_forward_activations_except(self, keras_mod):
234241
"""

0 commit comments

Comments
 (0)