I ran the code on tf 2.4 and and face a different error on nightly, please find the gist here

As I mentioned, I don't think tf-nightly-gpu is running on the GPU in Colab. The issue is only present on GPU. In the gist you sent, the cell where it searches for the GPU returns:

Tensorflow version ==  2.5.0-dev20210124

---------------------------------------------------------------------------
SystemError                               Traceback (most recent call last)
<ipython-input-8-6f54130e32b7> in <module>()
      4 print(device_name)
      5 if device_name != '/device:GPU:0':
----> 6   raise SystemError('GPU device not found')
      7 print('Found GPU at: {}'.format(device_name))

SystemError: GPU device not found

Locally, where it is running on GPU, TF nightly fails with the same error as TF2.4, so this issue is still present in nightly, at least on my local pip install tf-nightly-gpu installation.

@djoshea this looks like a reasonable request, will you be able to send a PR for this?

I could potentially try developing it, but I unfortunately don't know where to start. Is there a guide somewhere to implementing new ops for the GPU?

I could potentially try developing it, but I unfortunately don't know where to start. Is there a guide somewhere to implementing new ops for the GPU?

I don't think we have a detailed guide for this, but if you're comfortable with C++ you could try to start at the check that's failing.

It would be nice if this could be implemented soon. Using map_fn with RaggedTensors is very convenient when working with data of different shapes. Unfortunately one must run these operations on CPU right now.

I'm relying on map_fn to extract ragged image patches using bounding boxes so I can then produce ragged bounding-box centered feature maps via convolutions. I am encountering the same bug with the gradient step inside a tf.function (I can successfully run outside tf.function). The ragged image patches method is enabling a large speed up with a reduced memory footprint so it would be very useful to enable this in graph mode on GPUs.

For me RaggedTensors + map_fn are quite an enabler because I am implementing some kind of link prediction in a graph neural network. Having TensorFlow handling the individual samples allows for pretty safe programming as I do not have to put all samples in a batch into a large, disconnected graph and make sure that no nodes of different graphs will be connected. Using TensorFlow this way results in nice code that is fast to write in contrast to solutions with masks which may need to be changed every time you change your computation. Having this in GPU would be awesome. 💯

I would love to see this fixed, especially as iterating over RaggedTensors seems to not properly work in a lot of cases.

I ran the code in TF v2.5 and face error ,please find the gist here..Thanks !

Hi @sanjoy any updates on this? This is a really needed feature.

Hi, i will also very appreciate if this bug get fixed.
I think it's somehow related to No gradient defined for operation RaggedTensorFromVariant and No gradient defined for operation RaggedTensorFromVariant / or no gradients at all
which are fixed with commit be6b1fd by @edloper .

The part with REGISTER_UNARY_VARIANT_UNARY_OP_FUNCTION( ZEROS_LIKE_VARIANT_UNARY_OP, DEVICE_CPU, RaggedTensorVariant, RaggedTensorVariantZerosLike<CPUDevice>); just registers the zeros_like OP for CPU device but not for GPU.

Maybe this will be a quick win to fix.
@edloper ,it would be very kind, if you can have a look at this.
Thanks.

Hi, thanks for raising the issue. According to @edloper "Basically, RaggedTensorVariant objects should never be copied to GPU, because we can't do anything useful with them there. But Placer isn't currently smart enough to figure that out (it just sees a Variant tensor, and doesn't know what kind of value it contains)." We have a project going on right now that hopefully will fix the issue.

Hello. Could you please inform if there has been any progress on the issue?

For me, it resulted in a sudden sixfold increase of the train step duration due to having to move to CPU (compared to residing to a contrived solution with non-ragged tensors)

Also needing this resolved

I faced the same issue and resolved it by rolling my own map_fn via TensorArray. Here's a rough sketch:

def map_fn2(fn, elems, fn_output_signature):
    batch_size = tf.shape(tf.nest.flatten(elems)[0])[0]
    arr = tf.TensorArray(
        fn_output_signature.dtype, size=batch_size, element_shape=fn_output_signature.shape)
    for i in tf.range(batch_size):
        arr = arr.write(i, fn(tf.nest.map_structure(lambda x: x[i], elems)))
    return arr.stack()

As of TF v. 2.15, still not resolved.
It would be nice if the limitation on using RaggedTensors with map_fn on GPU will be mentioned on map_fn documentation.

Agreed, I would like to see this resolved as it is pretty important functionality for using map_fn and ragged tensors. @sanjoy @djoshea