Skip to content

Commit

Permalink
[Bugfix][Frontend][Keras]Fix a corner case bug in softmax converter o…
Browse files Browse the repository at this point in the history
…f keras frontend (#15337)

* Fix softmax converter about keras

* add new test cases to capture the bug

* Update keras.py
  • Loading branch information
jikechao authored Jul 18, 2023
1 parent c4f10cd commit e2d6511
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 2 deletions.
6 changes: 4 additions & 2 deletions python/tvm/relay/frontend/keras.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,11 +131,13 @@ def _convert_advanced_activation(inexpr, keras_layer, etab, data_layout, input_s

if act_type == "Softmax":
axis = keras_layer.axis
dims = len(input_shape)
dims = len(input_shape) if input_shape else 0
if isinstance(axis, list):
raise tvm.error.OpAttributeUnImplemented(f"Softmax with axes {axis} is not supported.")
if data_layout == "NCHW":
if axis == -1:
if dims == 0:
axis = 0
elif axis == -1:
axis = 1
else:
axis = axis + 1 if axis < dims - 1 else 1
Expand Down
7 changes: 7 additions & 0 deletions tests/python/frontend/keras/test_forward.py
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,13 @@ def test_forward_activations(self, keras_mod):
keras_model = keras_mod.models.Model(data, x)
verify_keras_frontend(keras_model)
verify_keras_frontend(keras_model, need_transpose=False, layout="NHWC")
# Test the input dimension = 1
data = keras_mod.layers.Input(shape=(11,))
act_func = keras_mod.layers.Softmax()
x = act_func(data)
keras_model = keras_mod.models.Model(data, x)
verify_keras_frontend(keras_model)
verify_keras_frontend(keras_model, need_transpose=False, layout="NHWC")

def test_forward_activations_except(self, keras_mod):
"""
Expand Down

0 comments on commit e2d6511

Please sign in to comment.