# Implement Softmax() Function in NumPy – NumPy Tutorial

By | August 12, 2021

In this tutorial, we will use an example to show you how to implement softmax function using numpy. You use code directly.

## Softmax function

Softmax function is defined as:

In numpy, if we compute softmax value of an array, we may get underflow and overflow problem. Here is a tutorial:

Implement Softmax Function Without Underflow and Overflow Problem – Deep Learning Tutorial

## How to implement softmax function for 1D and 2D array in numpy?

Look at example code:

def softmax(x):
x_1d = False
if x.ndim == 1:
x = np.expand_dims(x, axis = 0)
x_1d = True
x = x - np.max(x, axis=1, keepdims=True)
x = np.exp(x)
x = x / np.sum(x, axis=1, keepdims=True)
if x_1d:
x = np.squeeze(x)
return x

This function supports 1D and 2D numpy array.

We can test this function as follows:

x = np.array([[1,2,3],[4,5,6]])
print(softmax(x))

x = np.array([1,2,3])
print(softmax(x))

Run this code, you will get:

[[0.09003057 0.24472847 0.66524096]
[0.09003057 0.24472847 0.66524096]]
[0.09003057 0.24472847 0.66524096]