Re order and tidy up
-from typing import Any
-
import numpy as np
+from numpy import ndarray
from py4j.java_gateway import java_import, JavaObject, JavaMember, is_instance_of
+from typing import Any
from keanu.base import JavaObjectWrapper
from keanu.context import KeanuContext
@@ -73,25 +73,24 @@ class Tensor(JavaObjectWrapper):
raise NotImplementedError("Generic types in a ndarray are not supported. Was given {}".format(type(scalar)))
@staticmethod
- def _ndarray_using_bytestream(java_tensor):
+ def _to_scalar_or_ndarray(java_tensor: JavaObject, return_as_primitive: bool = False) -> Any:
+ if java_tensor.getRank() == 0:
+ scalar_array = np.array(java_tensor.scalar())
+ return scalar_array.item() if return_as_primitive else scalar_array
+ else:
+ return Tensor.__get_ndarray_from_tensor(java_tensor).reshape(java_tensor.getShape())
+
+ @staticmethod
+ def __get_ndarray_from_tensor(java_tensor) -> ndarray:
+ # Performance is much better using byte arrays where possible.
+ # https://stackoverflow.com/questions/39095994/fast-conversion-of-java-array-to-numpy-array-py4j
if is_instance_of(k._gateway, java_tensor, "io.improbable.keanu.tensor.dbl.DoubleTensor"):
byteArray = k.jvm_view().Py4jUtils.toByteArray(java_tensor.asFlatDoubleArray())
doubleArray = np.frombuffer(byteArray, np.float64)
return doubleArray
elif is_instance_of(k._gateway, java_tensor, "io.improbable.keanu.tensor.intgr.IntegerTensor"):
byteArray = k.jvm_view().Py4jUtils.toByteArray(java_tensor.asFlatIntegerArray())
- doubleArray = np.frombuffer(byteArray, np.int32)
- return doubleArray
+ intArray = np.frombuffer(byteArray, np.int32)
+ return intArray
else:
return np.array(list(java_tensor.asFlatArray()))
-
-
- @staticmethod
- def _to_scalar_or_ndarray(java_tensor: JavaObject, return_as_primitive: bool = False) -> Any:
- if java_tensor.getRank() == 0:
- scalar_array = np.array(java_tensor.scalar())
- return scalar_array.item() if return_as_primitive else scalar_array
- else:
- return Tensor._ndarray_using_bytestream(java_tensor).reshape(java_tensor.getShape())
- # return np.array(list(java_tensor.asFlatArray())).reshape(java_tensor.getShape())
-
