In the first plot of the Square and ReLU functions:
ax[0].set_title("Square function")
ax[0].set_xlabel("input")
ax[0].set_ylabel("input")
ylabel is a duplicate of xlabel.
should have been:
ax[0].set_title("Square function")
ax[0].set_xlabel("input")
ax[0].set_ylabel("output")