A search for resonant and non-resonant pair production of Higgs bosons in the final state is presented. The search uses 36.1 fb of collision data with TeV recorded by the ATLAS experiment at the LHC in 2015 and 2016. The semileptonic and fully hadronic decays of the -lepton pair are considered. No significant excess above the expected background is observed in the data. The cross-section times branching ratio for non-resonant Higgs boson pair production is constrained to be less than 30.9 fb, 12.7 times the Standard Model expectation, at 95% confidence level. The data are also analyzed to probe resonant Higgs boson pair production, constraining a model with an extended Higgs sector based on two doublets and a Randall-Sundrum bulk graviton model. Upper limits are placed on the resonant Higgs boson pair production cross-section times branching ratio, excluding resonances in the mass range 305~{\rm GeV} < m_X < 402~{\rm GeV} in the simplified hMSSM minimal supersymmetric model for and excluding bulk Randall-Sundrum gravitons in the mass range 325~{\rm GeV} < m_{G_{\mathrm{KK}}} < 885~{\rm GeV} for .