[updated] improve softmax implementation
diff --git a/keras/activations.py b/keras/activations.py
index a0ae77385fe..cf0f2134866 100644
--- a/keras/activations.py
+++ b/keras/activations.py
@@ -25,7 +25,9 @@ def softmax(x, axis=-1):
         ValueError: In case `dim(x) == 1`.
     """
     ndim = K.ndim(x)
-    if ndim == 2:
+    if ndim == 1:
+        raise ValueError('Cannot apply softmax to a tensor that is 1D')
+    elif ndim == 2:
         return K.softmax(x)
     elif ndim > 2:
         e = K.exp(x - K.max(x, axis=axis, keepdims=True))
diff --git a/keras/backend/theano_backend.py b/keras/backend/theano_backend.py
index 17ce97fb13f..9a47923c058 100644
--- a/keras/backend/theano_backend.py
+++ b/keras/backend/theano_backend.py
@@ -1716,10 +1716,11 @@ def relu(x, alpha=0., max_value=None):
 
 
 def softmax(x, axis=-1):
-    if axis == -1 or axis == x.ndim - 1:
+    if (axis == -1 or axis == x.ndim - 1) and x.ndim == 2:
         return T.nnet.softmax(x)
-    return T.exp(x - x.max()) / T.exp(
-        x - x.max()).sum(axis=axis, keepdims=True)
+    xm = x.max(axis=axis, keepdims=True)
+    return T.exp(x - xm) / T.exp(
+        x - xm).sum(axis=axis, keepdims=True)
 
 
 def softplus(x):
diff --git a/tests/keras/activations_test.py b/tests/keras/activations_test.py
index 5be59def8c8..fe93075fa01 100644
--- a/tests/keras/activations_test.py
+++ b/tests/keras/activations_test.py
@@ -79,6 +79,23 @@ def test_softmax_invalid():
         f = K.function([x], [activations.softmax(x)])
 
 
+def test_softmax_3d():
+    """Test using a reference implementation of softmax.
+    """
+    def softmax(values, axis):
+        m = np.max(values, axis=axis, keepdims=True)
+        e = np.exp(values - m)
+        return e / np.sum(e, axis=axis, keepdims=True)
+
+    x = K.placeholder(ndim=3)
+    f = K.function([x], [activations.softmax(x, axis=1)])
+    test_values = get_standard_values()[:, :, np.newaxis].copy()
+
+    result = f([test_values])[0]
+    expected = softmax(test_values, axis=1)
+    assert_allclose(result, expected, rtol=1e-05)
+
+
 def test_time_distributed_softmax():
     x = K.placeholder(shape=(1, 1, 5))
     f = K.function([x], [activations.softmax(x)])
diff --git a/tests/keras/backend/backend_test.py b/tests/keras/backend/backend_test.py
index f1d44c1a1de..eb614eda27c 100644
--- a/tests/keras/backend/backend_test.py
+++ b/tests/keras/backend/backend_test.py
@@ -975,6 +975,7 @@ def test_nn_operations(self):
         check_single_tensor_operation('tanh', (4, 2), WITH_NP)
 
         check_single_tensor_operation('softmax', (4, 10), WITH_NP)
+        check_single_tensor_operation('softmax', (4, 5, 3), WITH_NP, axis=1)
         check_single_tensor_operation('softmax', (4, 5, 3, 10), WITH_NP, axis=2)
 
         check_two_tensor_operation('binary_crossentropy', (4, 2), (4, 2), WITH_NP, from_logits=True)

### Summary

This PR is an update of #10027 which I kinda broke. Sorry about it.

### Related Issues

### PR Overview

I moved the correctness test in the backend tests. I think it belongs there rather than in the activations tests.

- [ ] This PR requires new unit tests [y/n] (make sure tests are included)
- [ ] This PR requires to update the documentation [y/n] (make sure the docs are up-to-date)
- [x] This PR is backwards compatible [y/n]
- [ ] This PR changes the current API [y/n] (all API changes need to be approved by fchollet)
