mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
make sure swig_ptr and rev_swig_ptr work on all primitive types (#1382)
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/1382 The array types supported for swig_ptr were not complete. This diff fixes that. Reviewed By: wickedfoo Differential Revision: D23411297 fbshipit-source-id: d94249b140aeb8c8179d9da3fcbc97eb034eac91
This commit is contained in:
committed by
Facebook GitHub Bot
parent
8dfd51d3a7
commit
c97f890651
@@ -210,6 +210,7 @@ namespace std {
|
||||
%include <std_pair.i>
|
||||
%include <std_map.i>
|
||||
|
||||
// primitive array types
|
||||
%template(FloatVector) std::vector<float>;
|
||||
%template(DoubleVector) std::vector<double>;
|
||||
%template(ByteVector) std::vector<uint8_t>;
|
||||
@@ -220,6 +221,10 @@ namespace std {
|
||||
%template(LongVector) std::vector<long>;
|
||||
%template(LongLongVector) std::vector<long long>;
|
||||
%template(IntVector) std::vector<int>;
|
||||
%template(UInt32Vector) std::vector<uint32_t>;
|
||||
%template(Int16Vector) std::vector<int16_t>;
|
||||
%template(UInt16Vector) std::vector<uint16_t>;
|
||||
|
||||
%template(FloatVectorVector) std::vector<std::vector<float> >;
|
||||
%template(ByteVectorVector) std::vector<std::vector<unsigned char> >;
|
||||
%template(LongVectorVector) std::vector<std::vector<long> >;
|
||||
@@ -799,15 +804,24 @@ PyObject *swig_ptr (PyObject *a)
|
||||
if(PyArray_TYPE(ao) == NPY_FLOAT64) {
|
||||
return SWIG_NewPointerObj(data, SWIGTYPE_p_double, 0);
|
||||
}
|
||||
if(PyArray_TYPE(ao) == NPY_INT32) {
|
||||
return SWIG_NewPointerObj(data, SWIGTYPE_p_int, 0);
|
||||
}
|
||||
if(PyArray_TYPE(ao) == NPY_UINT8) {
|
||||
return SWIG_NewPointerObj(data, SWIGTYPE_p_unsigned_char, 0);
|
||||
}
|
||||
if(PyArray_TYPE(ao) == NPY_INT8) {
|
||||
return SWIG_NewPointerObj(data, SWIGTYPE_p_char, 0);
|
||||
}
|
||||
if(PyArray_TYPE(ao) == NPY_UINT16) {
|
||||
return SWIG_NewPointerObj(data, SWIGTYPE_p_unsigned_short, 0);
|
||||
}
|
||||
if(PyArray_TYPE(ao) == NPY_INT16) {
|
||||
return SWIG_NewPointerObj(data, SWIGTYPE_p_short, 0);
|
||||
}
|
||||
if(PyArray_TYPE(ao) == NPY_UINT32) {
|
||||
return SWIG_NewPointerObj(data, SWIGTYPE_p_unsigned_int, 0);
|
||||
}
|
||||
if(PyArray_TYPE(ao) == NPY_INT32) {
|
||||
return SWIG_NewPointerObj(data, SWIGTYPE_p_int, 0);
|
||||
}
|
||||
if(PyArray_TYPE(ao) == NPY_UINT64) {
|
||||
#ifdef SWIGWORDSIZE64
|
||||
return SWIG_NewPointerObj(data, SWIGTYPE_p_unsigned_long, 0);
|
||||
@@ -870,8 +884,13 @@ PyObject * rev_swig_ptr(ctype *src, size_t size);
|
||||
%enddef
|
||||
|
||||
REV_SWIG_PTR(float, NPY_FLOAT32);
|
||||
REV_SWIG_PTR(int, NPY_INT32);
|
||||
REV_SWIG_PTR(double, NPY_FLOAT64);
|
||||
REV_SWIG_PTR(unsigned char, NPY_UINT8);
|
||||
REV_SWIG_PTR(char, NPY_INT8);
|
||||
REV_SWIG_PTR(unsigned short, NPY_UINT16);
|
||||
REV_SWIG_PTR(short, NPY_INT16);
|
||||
REV_SWIG_PTR(int, NPY_INT32);
|
||||
REV_SWIG_PTR(unsigned int, NPY_UINT32);
|
||||
REV_SWIG_PTR(int64_t, NPY_INT64);
|
||||
REV_SWIG_PTR(uint64_t, NPY_UINT64);
|
||||
|
||||
@@ -1038,4 +1057,3 @@ struct MapLong2Long {
|
||||
%}
|
||||
|
||||
// End of file...
|
||||
|
||||
|
||||
@@ -628,5 +628,25 @@ class TestSWIGWrap(unittest.TestCase):
|
||||
|
||||
assert isinstance(index2, faiss.IndexRefineFlat)
|
||||
|
||||
def do_test_array_type(self, dtype):
|
||||
""" tests swig_ptr and rev_swig_ptr for this type of array """
|
||||
a = np.arange(12).astype(dtype)
|
||||
ptr = faiss.swig_ptr(a)
|
||||
print(ptr)
|
||||
a2 = faiss.rev_swig_ptr(ptr, 12)
|
||||
np.testing.assert_array_equal(a, a2)
|
||||
|
||||
def test_all_array_types(self):
|
||||
self.do_test_array_type('float32')
|
||||
self.do_test_array_type('float64')
|
||||
self.do_test_array_type('int8')
|
||||
self.do_test_array_type('uint8')
|
||||
self.do_test_array_type('int16')
|
||||
self.do_test_array_type('uint16')
|
||||
self.do_test_array_type('int32')
|
||||
self.do_test_array_type('uint32')
|
||||
self.do_test_array_type('int64')
|
||||
self.do_test_array_type('uint64')
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user