Function half_bf16_works

tests/array.rs:605–644

()

#[cfg(feature = "half")]
#[test]
fn half_bf16_works() {
    Python::attach(|py| {
        let np = py.eval(c_str!("__import__('numpy')"), None, None).unwrap();
        // NumPy itself does not provide a `bfloat16` dtype itself,
        // so we import ml_dtypes which does register such a dtype.
        let mldt = py
            .eval(c_str!("__import__('ml_dtypes')"), None, None)
            .unwrap();
        let locals = [("np", &np), ("mldt", &mldt)].into_py_dict(py).unwrap();

        let array = py
            .eval(
                c_str!("np.array([[1, 2], [3, 4]], dtype='bfloat16')"),
                None,
                Some(&locals),
            )
            .unwrap()
            .cast_into::<PyArray2<bf16>>()
            .unwrap();

        assert_eq!(
            array.readonly().as_array(),
            array![
                [bf16::from_f32(1.0), bf16::from_f32(2.0)],
                [bf16::from_f32(3.0), bf16::from_f32(4.0)]
            ]
        );

        array
            .readwrite()
            .as_array_mut()
            .map_inplace(|value| *value *= bf16::from_f32(2.0));

        py_run!(
            py,
            array np,
            "assert np.all(array == np.array([[2, 4], [6, 8]], dtype='bfloat16'))"
        );
    });
}

#[test]
fn ascii_strings_with_explicit_dtype_works() {