diff --git a/src/awkward/operations/ak_from_raggedtensor.py b/src/awkward/operations/ak_from_raggedtensor.py index 1c895506c2..69c3558281 100644 --- a/src/awkward/operations/ak_from_raggedtensor.py +++ b/src/awkward/operations/ak_from_raggedtensor.py @@ -28,20 +28,43 @@ def from_raggedtensor(array): def _impl(array): + try: + import tensorflow as tf + except ImportError as err: + raise ImportError( + """to use ak.from_raggedtensor, you must install the 'tensorflow' package with: + + pip install tensorflow +or + conda install tensorflow""" + ) from err + try: # get the flat values - content = array.flat_values.numpy() + content = array.flat_values except AttributeError as err: raise TypeError( """only RaggedTensor can be converted to awkward array""" ) from err - # convert them to ak.contents right away + + # handle gpu and cpu instances separately + device = content.device + + # since TensorFlow currently does not support + # int32 variables being placed on the GPU, use CPU for them instead + if content.dtype == tf.int32: + device = "cpu" + + content = _tensor_to_np_or_cp(content, device) + + # convert flat_values to ak.contents right away content = ak.contents.NumpyArray(content) # get the offsets offsets_arr = [] for splits in array.nested_row_splits: - split = splits.numpy() + # handle gpu and cpu instances separately + split = _tensor_to_np_or_cp(splits, device) # convert to ak.index offset = ak.index.Index64(split) offsets_arr.append(offset) @@ -55,6 +78,18 @@ def _impl(array): return ak.Array(_recursive_call(content, offsets_arr, 0)) +def _tensor_to_np_or_cp(array, device): + import tensorflow as tf + + if "GPU" in device: + from awkward._nplikes.cupy import Cupy + + cp = Cupy.instance() + return cp.from_dlpack(tf.experimental.dlpack.to_dlpack(array)) + else: + return array.numpy() + + def _recursive_call(content, offsets_arr, count): if count == len(offsets_arr) - 2: return ak.contents.ListOffsetArray(