Merge pull request #314 from PyO3/half

Add support for 16-bit floating point numbers.

Author: adamreichold

.github/workflows/ci.yml

@@ -121,6 +121,7 @@ jobs:
           import toml
           cargo_toml = toml.load("Cargo.toml")
           cargo_toml["dependencies"]["ndarray"] = "0.13.1"
+          cargo_toml["dependencies"]["parking_lot"] = "0.11.2"
           cargo_toml["dependencies"]["num-complex"] = "0.2.4"
           cargo_toml["workspace"] = {}
           with open("Cargo.toml", "w") as f:
@@ -130,15 +131,17 @@ jobs:
       - name: Generate lockfile
         run: cargo generate-lockfile
         working-directory: examples/simple
-      - name: Unify dependencies on ndarray to 0.13.1
+      - name: Unify dependencies on ndarray to 0.13.1 and on parking_lot to 0.11.2
         run: |
           import toml
           import subprocess
           cargo_lock = toml.load("Cargo.lock")
           for pkg in cargo_lock["package"]:
+            pkg_id = pkg["name"] + ":" + pkg["version"]
             if pkg["name"] == "ndarray" and pkg["version"] != "0.13.1":
-              pkg_id = pkg["name"] + ":" + pkg["version"]
               subprocess.run(["cargo", "update", "--package", pkg_id, "--precise", "0.13.1"], check=True)
+            if pkg["name"] == "parking_lot" and pkg["version"] != "0.11.2":
+              subprocess.run(["cargo", "update", "--package", pkg_id, "--precise", "0.11.2"], check=True)
         working-directory: examples/simple
         shell: python
       - name: Test example

CHANGELOG.md

@@ -4,6 +4,7 @@
   - Add dynamic borrow checking to safely construct references into the interior of NumPy arrays. ([#274](https://github.com/PyO3/rust-numpy/pull/274))
     - The deprecated iterator builders `NpySingleIterBuilder::{readonly,readwrite}` and `NpyMultiIterBuilder::add_{readonly,readwrite}` now take referencces to `PyReadonlyArray` and `PyReadwriteArray` instead of consuming them.
     - The destructive `PyArray::resize` method is now unsafe if used without an instance of `PyReadwriteArray`. ([#302](https://github.com/PyO3/rust-numpy/pull/302))
+  - Add support for IEEE 754-2008 16-bit floating point numbers via an optional dependency on the `half` crate. ([#314](https://github.com/PyO3/rust-numpy/pull/314))
   - The `inner`, `dot` and `einsum` functions can also return a scalar instead of a zero-dimensional array to match NumPy's types ([#285](https://github.com/PyO3/rust-numpy/pull/285))
   - Deprecate `PyArray::from_exact_iter` after optimizing `PyArray::from_iter`. ([#292](https://github.com/PyO3/rust-numpy/pull/292))
   - Fix returning invalid slices from `PyArray::{strides,shape}` for rank zero arrays. ([#303](https://github.com/PyO3/rust-numpy/pull/303))

Cargo.toml

@@ -16,6 +16,7 @@ license = "BSD-2-Clause"
 
 [dependencies]
 ahash = "0.7"
+half = { version = "1.8", default-features = false, optional = true }
 libc = "0.2"
 num-complex = ">= 0.2, < 0.5"
 num-integer = "0.1"
@@ -28,3 +29,6 @@ pyo3 = { version = "0.16", default-features = false, features = ["auto-initializ
 
 [workspace]
 members = ["examples/*"]
+
+[package.metadata.docs.rs]
+all-features = true

src/dtype.rs

@@ -4,6 +4,8 @@ use std::os::raw::{
 };
 use std::ptr;
 
+#[cfg(feature = "half")]
+use half::f16;
 use num_traits::{Bounded, Zero};
 use pyo3::{
     exceptions::{PyIndexError, PyValueError},
@@ -471,6 +473,9 @@ impl_element_scalar!(u8, u16, u32, u64);
 impl_element_scalar!(f32 => NPY_FLOAT);
 impl_element_scalar!(f64 => NPY_DOUBLE);
 
+#[cfg(feature = "half")]
+impl_element_scalar!(f16 => NPY_HALF);
+
 impl_element_scalar!(Complex32 => NPY_CFLOAT,
     #[doc = "Complex type with `f32` components which maps to `numpy.csingle` (`numpy.complex64`)."]);
 impl_element_scalar!(Complex64 => NPY_CDOUBLE,

tests/array.rs

@@ -1,5 +1,7 @@
 use std::mem::size_of;
 
+#[cfg(feature = "half")]
+use half::f16;
 use ndarray::{array, s, Array1, Dim};
 use numpy::{
     dtype, get_array_module, pyarray, PyArray, PyArray1, PyArray2, PyArrayDescr, PyArrayDyn,
@@ -505,3 +507,41 @@ fn get_works() {
         }
     });
 }
+
+#[cfg(feature = "half")]
+#[test]
+fn half_works() {
+    Python::with_gil(|py| {
+        let np = py.eval("__import__('numpy')", None, None).unwrap();
+        let locals = [("np", np)].into_py_dict(py);
+
+        let array = py
+            .eval(
+                "np.array([[1, 2], [3, 4]], dtype='float16')",
+                None,
+                Some(locals),
+            )
+            .unwrap()
+            .downcast::<PyArray2<f16>>()
+            .unwrap();
+
+        assert_eq!(
+            array.readonly().as_array(),
+            array![
+                [f16::from_f32(1.0), f16::from_f32(2.0)],
+                [f16::from_f32(3.0), f16::from_f32(4.0)]
+            ]
+        );
+
+        array
+            .readwrite()
+            .as_array_mut()
+            .map_inplace(|value| *value *= f16::from_f32(2.0));
+
+        py_run!(
+            py,
+            array np,
+            "np.testing.assert_array_almost_equal(array, np.array([[2, 4], [6, 8]], dtype='float16'))"
+        );
+    });
+}

x.py

@@ -28,46 +28,46 @@ def default(args):
     run("cargo", "fmt")
 
     if nightly():
-        run("cargo", "clippy", "--workspace", "--tests", "--benches")
+        run("cargo", "clippy", "--workspace", "--all-features", "--tests", "--benches")
     else:
-        run("cargo", "clippy", "--workspace", "--tests")
+        run("cargo", "clippy", "--workspace", "--all-features", "--tests")
 
-    run("cargo", "test", "--lib", "--tests")
+    run("cargo", "test", "--all-features", "--lib", "--tests")
 
 
 def check(args):
     run("cargo", "fmt", "--", "--check")
 
-    run("cargo", "clippy", "--workspace", "--tests", "--", "--deny", "warnings")
+    run("cargo", "clippy", "--workspace", "--all-features", "--tests", "--", "--deny", "warnings")
 
 
 def doc(args):
     if args.name is None:
-        run("cargo", "test", "--doc")
+        run("cargo", "test", "--all-features", "--doc")
     else:
-        run("cargo", "test", "--doc", args.name)
+        run("cargo", "test", "--all-features", "--doc", args.name)
 
     if args.open:
-        run("cargo", "doc", "--open")
+        run("cargo", "doc", "--all-features", "--open")
     else:
-        run("cargo", "doc")
+        run("cargo", "doc", "--all-features")
 
 
 def test(args):
     if args.name is None:
-        run("cargo", "test", "--tests")
+        run("cargo", "test", "--all-features", "--tests")
     else:
-        run("cargo", "test", "--test", args.name)
+        run("cargo", "test", "--all-features", "--test", args.name)
 
 
 def bench(args):
     if not nightly():
         sys.exit("Benchmarks require a nightly build of the Rust compiler.")
 
     if args.name is None:
-        run("cargo", "bench", "--benches")
+        run("cargo", "bench", "--all-features", "--benches")
     else:
-        run("cargo", "bench", "--bench", args.name)
+        run("cargo", "bench", "--all-features", "--bench", args.name)
 
 
 def examples(args):