norm_sq with axis returns 1st element of vector instead of vector

[nbecker]

This code for norm_sq should return a vector, but instead returns a scalar: the 1st element of the vector

-----bug.cc

#define FORCE_IMPORT_ARRAY
#define _GNU_SOURCE
#define HAVE_CBLAS 1

#include <cmath>
#include <xtensor/xmath.hpp>
#include <xtensor/xtensor.hpp>
#include <xtensor/xarray.hpp>
#include "xtensor-python/pytensor.hpp"
#include "xtensor-python/pyarray.hpp"
#include "pybind11/pybind11.h"
#include <xtensor/xtensor_simd.hpp>
//#include <xtensor/xexpression.hpp>
#include <xtensor/xcomplex.hpp>
#include "xtensor/xnorm.hpp"
//#include "xtensor-python/pyvectorize.hpp"
//#include "xtensor-blas/xlinalg.hpp"

namespace py = pybind11;

PYBIND11_MODULE (bug, m) {
  xt::import_numpy();

    m.def ("norm_sq", [](xt::pyarray<double> const& x) { return xt::norm_sq (x)(); },
	 py::arg("in").noconvert()
	 );

  m.def ("norm_sq", [](xt::pyarray<double> const& x, std::vector<size_t>const& ax) { return xt::norm_sq (x, ax, xt::evaluation_strategy::immediate)(); },
	 py::arg("in").noconvert(),
	 py::arg("axes")
	 );

  m.def ("norm_sq", [](xt::pyarray<double> const& x, size_t ax) {
    auto shape = std::vector<size_t>{ax};
    return xt::norm_sq (x, shape, xt::evaluation_strategy::immediate)(); },
    py::arg("in").noconvert(),
    py::arg("axes")
    );

  m.def ("norm_sq", [](xt::pyarray<std::complex<double>> const& x) { return xt::norm_sq (x)(); },
	 py::arg("in").noconvert()
	 );

  m.def ("norm_sq", [](xt::pyarray<std::complex<double>> const& x, std::vector<size_t>const& ax) { return xt::norm_sq (x, ax, xt::evaluation_strategy::immediate)(); },
	 py::arg("in").noconvert(),
	 py::arg("axes")
	 );

  m.def ("norm_sq", [](xt::pyarray<std::complex<double>> const& x, size_t ax) {
    auto shape = std::vector<size_t>{ax};
    return xt::norm_sq (x, shape, xt::evaluation_strategy::immediate)(); },
    py::arg("in").noconvert(),
    py::arg("axes")
    );
}

---

test_bug.py

import numpy as np
from bug import norm_sq

u = np.array ([[1., 2., 3.],
               [4., 5., 6.],
               [7., 8., 9.]])
print (norm_sq (u, axes=0))

print (np.linalg.norm(u, axis=0)**2)
-------------

In [65]: 66.0 << result of norm_sq(u,axes=0) [ 66. 93. 126.] << result of np.linalg.norm(u,axis=0)**2