add cuda backend support for ak.from_raggedtensor

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(