Skip to content

Commit

Permalink
Merge pull request #29 from datasig-ac-uk/fix/from-jax-array-issues
Browse files Browse the repository at this point in the history 
Fix/from jax array issues
  • Loading branch information
@inakleinbottle
inakleinbottle authored Aug 7, 2023
2 parents aa3776c + 2db3da5 commit d756bf2
Show file tree
Hide file tree
Showing 4 changed files with 179 additions and 46 deletions.
44 changes: 31 additions & 13 deletions roughpy/src/scalars/scalars.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -217,6 +217,10 @@ static bool try_fill_buffer_dlpack(
}

auto& dltensor = tensor->dl_tensor;

RPY_CHECK(dltensor.device.device_type == kDLCPU,
"no device support is currently available");

auto* data = reinterpret_cast<char*>(dltensor.data);
auto ndim = dltensor.ndim;
auto* shape = dltensor.shape;
Expand All @@ -243,12 +247,8 @@ static bool try_fill_buffer_dlpack(
for (auto i = 0; i < ndim; ++i) { size *= static_cast<idimn_t>(shape[i]); }

if (strides == nullptr) {
if (options.type == tensor_stype) {
buffer = scalars::ScalarArray({options.type, data}, size);
} else {
buffer.allocate_scalars(size);
options.type->convert_copy(buffer, {tensor_stype, data}, size);
}
buffer.allocate_scalars(size);
options.type->convert_copy(buffer, {tensor_stype, data}, size);
} else {
buffer.allocate_scalars(size);
scalars::ScalarPointer p(tensor_stype, data);
Expand Down Expand Up @@ -536,7 +536,10 @@ scalars::KeyScalarArray python::py_to_buffer(
update_dtype_and_allocate(result, options, 1, 0);

assign_py_object_to_scalar(result, object);
} else if (is_kv_pair(object, options.alternative_key)) {
return result;
}

if (is_kv_pair(object, options.alternative_key)) {
/*
* Now for tuples of length 2, which we expect to be a kv-pair
*/
Expand All @@ -547,12 +550,16 @@ scalars::KeyScalarArray python::py_to_buffer(
update_dtype_and_allocate(result, options, 1, 1);

handle_sequence_tuple(result, result.keys(), object, options);
} else if (py::hasattr(object, "__dlpack__")) {
return result;
}

if (py::hasattr(object, "__dlpack__")) {
// If we used the dlpack interface, then the result is
// already constructed.
try_fill_buffer_dlpack(result, options, object);
if (try_fill_buffer_dlpack(result, options, object)) { return result; }
}

} else if (py::isinstance<py::buffer>(object)) {
if (py::isinstance<py::buffer>(object)) {
// Fall back to the buffer protocol
auto info = py::reinterpret_borrow<py::buffer>(object).request();
auto type_id = py_buffer_to_type_id(info);
Expand All @@ -573,7 +580,10 @@ scalars::KeyScalarArray python::py_to_buffer(
options.shape.assign(info.shape.begin(), info.shape.end());
}

} else if (py::isinstance<py::dict>(object)) {
return result;
}

if (py::isinstance<py::dict>(object)) {
auto dict_arg = py::reinterpret_borrow<py::dict>(object);
options.shape.push_back(static_cast<idimn_t>(dict_arg.size()));

Expand All @@ -590,7 +600,10 @@ scalars::KeyScalarArray python::py_to_buffer(

handle_dict(ptr, key_ptr, options, dict_arg);
}
} else if (py::isinstance<py::sequence>(object)) {
return result;
}

if (py::isinstance<py::sequence>(object)) {
std::vector<py::object> leaves;
auto size_info = compute_size_and_type(options, leaves, object);

Expand Down Expand Up @@ -625,9 +638,14 @@ scalars::KeyScalarArray python::py_to_buffer(
}
}
}

return result;
}

return result;
RPY_THROW(
std::invalid_argument,
"could not parse argument to a valid scalar type"
);
}

scalars::Scalar
Expand Down
124 changes: 91 additions & 33 deletions scalars/src/scalar_type.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,21 +58,59 @@ const ScalarType* ScalarType::for_id(const string& id)
{
return ScalarType::of<double>();
}
const ScalarType* ScalarType::from_type_details(const BasicScalarInfo& details,
const ScalarDeviceInfo& device)
const ScalarType* ScalarType::from_type_details(
const BasicScalarInfo& details, const ScalarDeviceInfo& device
)
{
return ScalarType::of<double>();
RPY_CHECK(device.device_type == ScalarDeviceType::CPU);

switch (details.code) {
case ScalarTypeCode::Int: RPY_FALLTHROUGH;
case ScalarTypeCode::UInt: return nullptr;
case ScalarTypeCode::Float:
switch (details.bits) {
case 16: return ScalarType::of<half>();
case 32: return ScalarType::of<float>();
case 64: return ScalarType::of<double>();
default:
RPY_THROW(
std::invalid_argument,
"supported float bits are 16, 32, and 64 only"
);
}
case ScalarTypeCode::BFloat:
switch (details.bits) {
case 16: return ScalarType::of<bfloat16>();
default:
RPY_THROW(
std::invalid_argument,
"supported bfloat bits are 16 only"
);
}
case ScalarTypeCode::OpaqueHandle: RPY_FALLTHROUGH;
case ScalarTypeCode::Complex: RPY_FALLTHROUGH;
case ScalarTypeCode::Bool:
RPY_THROW(
std::invalid_argument,
"opaque handles, complex, and bool arguments are not "
"currently supported"
);
}
RPY_UNREACHABLE_RETURN(nullptr);
}
Scalar ScalarType::from(long long int numerator,
long long int denominator) const
Scalar
ScalarType::from(long long int numerator, long long int denominator) const
{
if (denominator == 0) { RPY_THROW(std::invalid_argument, "division by zero"); }
if (denominator == 0) {
RPY_THROW(std::invalid_argument, "division by zero");
}
auto result = allocate(1);
assign(result, numerator, denominator);
return {result, flags::OwnedPointer};
}
void ScalarType::convert_fill(ScalarPointer out, ScalarPointer in, dimn_t count,
const string& id) const
void ScalarType::convert_fill(
ScalarPointer out, ScalarPointer in, dimn_t count, const string& id
) const
{
if (!id.empty()) {
try {
Expand Down Expand Up @@ -193,6 +231,7 @@ static const pair<string, ScalarTypeInfo> reserved[]
#define U64IDX 3

static std::mutex s_scalar_type_cache_lock;

static std::unordered_map<string, const ScalarType*> s_scalar_type_cache {
{string("f32"), ScalarType::of<float>()},
{string("f64"), ScalarType::of<double>()},
Expand All @@ -211,15 +250,20 @@ void rpy::scalars::register_type(const ScalarType* type)
const auto& identifier = type->id();
for (const auto& i : reserved) {
if (identifier == i.first) {
RPY_THROW(std::runtime_error,"cannot register identifier " + identifier
+ ", it is reserved");
RPY_THROW(
std::runtime_error,
"cannot register identifier " + identifier
+ ", it is reserved"
);
}
}

auto& entry = s_scalar_type_cache[identifier];
if (entry != nullptr) {
RPY_THROW(std::runtime_error,"type with id " + identifier
+ " is already registered");
RPY_THROW(
std::runtime_error,
"type with id " + identifier + " is already registered"
);
}

entry = type;
Expand All @@ -237,8 +281,8 @@ const ScalarType* rpy::scalars::get_type(const string& id)
RPY_THROW(std::runtime_error, "no type with id " + id);
}

const ScalarType* scalars::get_type(const string& id,
const ScalarDeviceInfo& device)
const ScalarType*
scalars::get_type(const string& id, const ScalarDeviceInfo& device)
{
// TODO: Needs implementation
return nullptr;
Expand Down Expand Up @@ -298,8 +342,9 @@ ROUGHPY_MAKE_TYPE_ID_OF(rational_poly_scalar, "RationalPoly")
#undef ROUGHPY_MAKE_TYPE_ID_OF

#define MAKE_CONVERSION_FUNCTION(SRC, DST, SRC_T, DST_T) \
static void SRC##_to_##DST(ScalarPointer dst, ScalarPointer src, \
dimn_t count) \
static void SRC##_to_##DST( \
ScalarPointer dst, ScalarPointer src, dimn_t count \
) \
{ \
const auto* src_p = src.raw_cast<const SRC_T>(); \
auto* dst_p = dst.raw_cast<DST_T>(); \
Expand Down Expand Up @@ -340,33 +385,39 @@ static std::unordered_map<string, conversion_function> s_conversion_cache{

#undef ADD_DEF_CONV

static inline string type_ids_to_key(const string& src_type,
const string& dst_type)
static inline string
type_ids_to_key(const string& src_type, const string& dst_type)
{
return src_type + "->" + dst_type;
}

const conversion_function& rpy::scalars::get_conversion(const string& src_id,
const string& dst_id)
const conversion_function&
rpy::scalars::get_conversion(const string& src_id, const string& dst_id)
{
std::lock_guard<std::mutex> access(s_conversion_lock);

auto found = s_conversion_cache.find(type_ids_to_key(src_id, dst_id));
if (found != s_conversion_cache.end()) { return found->second; }

RPY_THROW(std::runtime_error,"no conversion function from " + src_id + " to "
+ dst_id);
RPY_THROW(
std::runtime_error,
"no conversion function from " + src_id + " to " + dst_id
);
}
void rpy::scalars::register_conversion(const string& src_id,
const string& dst_id,
conversion_function converter)
void rpy::scalars::register_conversion(
const string& src_id, const string& dst_id,
conversion_function converter
)
{
std::lock_guard<std::mutex> access(s_conversion_lock);

auto& found = s_conversion_cache[type_ids_to_key(src_id, dst_id)];
if (found != nullptr) {
RPY_THROW(std::runtime_error,"conversion from " + src_id + " to " + dst_id
+ " already registered");
RPY_THROW(
std::runtime_error,
"conversion from " + src_id + " to " + dst_id
+ " already registered"
);
} else {
found = std::move(converter);
}
Expand All @@ -380,13 +431,17 @@ ScalarType::ScalarType(ScalarTypeInfo info)
std::unique_ptr<RandomGenerator>
ScalarType::get_rng(const string& bit_generator, Slice<uint64_t> seed) const
{
RPY_THROW(std::runtime_error,
"no random number generators are defined for this scalar type");
RPY_THROW(
std::runtime_error,
"no random number generators are defined for this scalar type"
);
}

std::unique_ptr<BlasInterface> ScalarType::get_blas() const
{
RPY_THROW(std::runtime_error, "No blas/lapack available for this scalar type");
RPY_THROW(
std::runtime_error, "No blas/lapack available for this scalar type"
);
}

const std::string& rpy::scalars::id_from_basic_info(const BasicScalarInfo& info)
Expand All @@ -402,21 +457,24 @@ const std::string& rpy::scalars::id_from_basic_info(const BasicScalarInfo& info)
case 16: return reserved[I16IDX].second.id;
case 32: return reserved[I32IDX].second.id;
case 64: return reserved[I64IDX].second.id;
default: RPY_THROW(std::runtime_error, "unsupported integer type");
default:
RPY_THROW(std::runtime_error, "unsupported integer type");
}
case ScalarTypeCode::UInt:
switch (info.bits) {
case 8: return reserved[U8IDX].second.id;
case 16: return reserved[U16IDX].second.id;
case 32: return reserved[U32IDX].second.id;
case 64: return reserved[U64IDX].second.id;
default: RPY_THROW(std::runtime_error, "unsupported integer type");
default:
RPY_THROW(std::runtime_error, "unsupported integer type");
}
case ScalarTypeCode::Float:
switch (info.bits) {
case 32: return ScalarType::of<float>()->id();
case 64: return ScalarType::of<double>()->id();
default: RPY_THROW(std::runtime_error, "unsupported float type");
default:
RPY_THROW(std::runtime_error, "unsupported float type");
}
default: RPY_THROW(std::runtime_error, "unsupported scalar type");
}
Expand Down
Empty file added tests/scalars/dlpack/__init__.py
View file
Empty file.
57 changes: 57 additions & 0 deletions tests/scalars/dlpack/test_construct_from_jax.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@


import pytest

try:
import jax
from jax import numpy as jnp
except ImportError:
jax = None
jnp = None

import roughpy as rp
import numpy as np
from numpy.testing import assert_array_equal

@pytest.mark.skipif(jax is None, reason="JAX test require jax to be installed")
class TestJaxArrayCreation:

@pytest.fixture(scope="class")
def context(self):
return rp.get_context(2, 2, rp.SPReal)

@pytest.fixture(scope="class")
def prng_key(self):
return jax.random.PRNGKey(12345)

def test_increment_stream_from_jax_array(self):
array = jnp.array([
[-0.25860816, -0.36977386, 0.6619457, -0.50442713, 0.08028925, -1.06701028],
[-0.26208243, 0.22464547, 0.39521545, -0.62663144, -0.34344956, -1.67293704],
[-0.55824, -0.19376263, 0.86616075, -0.58314389, -0.69254208, -1.53291035],
[-0.52306908, -0.09234464, 1.17564034, -0.7388621, -0.91333717, -1.50844121],
[-0.80696738, -0.09417236, 0.75135314, -1.20548987, -1.42038512, -1.86834741],
[-0.6642682, -0.12166289, 1.04914618, -1.01415539, -1.58841276, -2.54356289]
])

stream = rp.LieIncrementStream.from_increments(np.array(array), width=6, depth=2, dtype="float32")

lsig01 = stream.log_signature(rp.RealInterval(0, 1))
lsig12 = stream.log_signature(rp.RealInterval(1, 2))
lsig24 = stream.log_signature(rp.RealInterval(2, 4))
lsig46 = stream.log_signature(rp.RealInterval(4, 6))

assert_array_equal(np.array(lsig01)[:6], array[0, :])
assert not np.any(np.isnan(lsig12))
assert not np.any(np.isnan(lsig24))
assert not np.any(np.isnan(lsig46))



@pytest.mark.xfail(condition=True, reason="No device support is currently available")
def test_create_tensor_from_jax_array(self, prng_key, context):
array = jax.random.uniform(prng_key, shape=(context.tensor_size(2),), dtype="float32", minval=-1.0, maxval=1.0)

ts = rp.FreeTensor(array, ctx=context)

assert_array_equal(np.array(ts), np.array(jax.device_get(array), copy=True))

0 comments on commit d756bf2

Please sign in to comment.