TypeError: softmax() got an unexpected keyword argument 'axis'

[nadamaguid]

I started getting this error suddently. Tensorflow has to be downgraded to 1.4 for the mnist.py example to work and it has been working fine. Suddenly, upon installing dist-keras recently I receive the following error:
Traceback (most recent call last): File "/home/hadoop/trial1200.py", line 277, in <module> model2.add(Activation('softmax')) File "/usr/local/lib64/python2.7/site-packages/keras/models.py", line 522, in add output_tensor = layer(self.outputs[0]) File "/usr/local/lib64/python2.7/site-packages/keras/engine/topology.py", line 619, in __call__ output = self.call(inputs, **kwargs) File "/usr/local/lib64/python2.7/site-packages/keras/layers/core.py", line 304, in call return self.activation(inputs) File "/usr/local/lib64/python2.7/site-packages/keras/activations.py", line 29, in softmax return K.softmax(x) File "/usr/local/lib64/python2.7/site-packages/keras/backend/tensorflow_backend.py", line 2963, in softmax return tf.nn.softmax(x, axis=axis) TypeError: softmax() got an unexpected keyword argument 'axis'
I double checked if using the --upgrade dist-keras command has installed anything differently and I noticed that the keras version was update from 2.1.5 to 2.1.6. Downgrading it back to 2.1.5 still gives the same error. Also if the version of tensorflow is upgraded, the transofrm function no longer works.
Is there a specific set up that needs to be done?