Awni Hannun
bc06cb9ff6
* pickle + dtype fix for numpy conversion * fix getattribute on Module base * remove unused function * fix tests * add topk to ops * fix doc
61 lines
1.8 KiB
C++
61 lines
1.8 KiB
C++
// Copyright © 2023-2024 Apple Inc.
|
|
#pragma once
|
|
|
|
// A patch to get float16_t to work with pybind11 numpy arrays
|
|
// Derived from:
|
|
// https://github.com/pybind/pybind11/issues/1776#issuecomment-492230679
|
|
|
|
#include <pybind11/numpy.h>
|
|
|
|
namespace pybind11::detail {
|
|
|
|
template <typename T>
|
|
struct npy_scalar_caster {
|
|
PYBIND11_TYPE_CASTER(T, _("PleaseOverride"));
|
|
using Array = array_t<T>;
|
|
|
|
bool load(handle src, bool convert) {
|
|
// Taken from Eigen casters. Permits either scalar dtype or scalar array.
|
|
handle type = dtype::of<T>().attr("type"); // Could make more efficient.
|
|
if (!convert && !isinstance<Array>(src) && !isinstance(src, type))
|
|
return false;
|
|
Array tmp = Array::ensure(src);
|
|
if (tmp && tmp.size() == 1 && tmp.ndim() == 0) {
|
|
this->value = *tmp.data();
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
static handle cast(T src, return_value_policy, handle) {
|
|
Array tmp({1});
|
|
tmp.mutable_at(0) = src;
|
|
tmp.resize({});
|
|
// You could also just return the array if you want a scalar array.
|
|
object scalar = tmp[tuple()];
|
|
return scalar.release();
|
|
}
|
|
};
|
|
|
|
// Similar to enums in `pybind11/numpy.h`. Determined by doing:
|
|
// python3 -c 'import numpy as np; print(np.dtype(np.float16).num)'
|
|
constexpr int NPY_FLOAT16 = 23;
|
|
|
|
// Kinda following:
|
|
// https://github.com/pybind/pybind11/blob/9bb3313162c0b856125e481ceece9d8faa567716/include/pybind11/numpy.h#L1000
|
|
template <>
|
|
struct npy_format_descriptor<float16_t> {
|
|
static constexpr auto name = _("float16");
|
|
static pybind11::dtype dtype() {
|
|
handle ptr = npy_api::get().PyArray_DescrFromType_(NPY_FLOAT16);
|
|
return reinterpret_borrow<pybind11::dtype>(ptr);
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct type_caster<float16_t> : npy_scalar_caster<float16_t> {
|
|
static constexpr auto name = _("float16");
|
|
};
|
|
|
|
} // namespace pybind11::detail
|