RaggedTensor support for sparse_categorical_crossentropy.

[pedro-r-marques]

Attempt to support RaggedTensors in sparse_categorical_crossentropy loss.

[foxik]

There is a typo - respectively; and I believe there should be a comma after respectively.

[pedro-r-marques]

Done.

[foxik]

Independently to the OP_REQUIRES issue, which I could not tackle, is the sparse_categorical_crossentropy expected to work in case there are multiple ragged dimensions? For example, the batch examples are images and we classify every pixel?

It seems that such a scenario works with binary_crossentropy from #47075, but if crashes here. Notably, the following seems to work:

sce_obj = tf.losses.SparseCategoricalCrossentropy()
y_true = tf.ragged.constant([[1, 1], [0]])
y_pred = tf.ragged.constant([[[0.1, 0.9], [0.1, 0.9]], [[0.9, 0.1]]])
print(sce_obj(y_true, y_pred))

but after adding another ragged dimension

y_true = tf.ragged.constant([[[1, 1]], [[0]]])
y_pred = tf.ragged.constant([[[[0.1, 0.9], [0.1, 0.9]]], [[[0.9, 0.1]]]]) # both without and with `inner_shape=(2,)`
print(sce_obj(y_true, y_pred))

it crashes with ValueError: TypeError: object of type 'RaggedTensor' has no len().

It seems to be connected to _convert_to_dense in _ragged_tensor_apply_loss, because if the _convert_to_dense just passes through the inputs unchanged, the example above works.

[pedro-r-marques]

@foxit Added a test with the example you provided and changed the approach a bit to avoid the warning generated by the previous code.

[foxik]

A gentle bump after two weeks -- we would be grateful for a review.

Thanks & cheers.

[tomerk]

Thank you!

[rthadur]

@pedro-r-marques can you please check build failures ?

[pedro-r-marques]

@rthadur Fixed. It was a pylint error for an unused import.

[pedro-r-marques]

@rthadur Can you please check why the CL doesn't built internally to google ?

[rthadur]

@pedro-r-marques seeing this error internally

tensorflow/python/util/dispatch.py", line 206, in wrapper
   return target(*args, **kwargs)
 File "/tensorflow/python/keras/losses.py", line 1736, in sparse_categorical_crossentropy
   y_pred = ops.convert_to_tensor_v2_with_dispatch(y_pred)
 File "/tensorflow/python/util/dispatch.py", line 206, in wrapper
   return target(*args, **kwargs)
 File "/tensorflow/python/framework/ops.py", line 1429, in convert_to_tensor_v2_with_dispatch
   value, dtype=dtype, dtype_hint=dtype_hint, name=name)
 File "/tensorflow/python/framework/ops.py", line 1439, in convert_to_tensor_v2
   as_ref=False)
 File "/tensorflow/python/profiler/trace.py", line 163, in wrapped
   return func(*args, **kwargs)
 File "/tensorflow/python/framework/ops.py", line 1564, in convert_to_tensor
   ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
 File "/tensorflow/python/framework/constant_op.py", line 340, in _constant_tensor_conversion_function
   return constant(v, dtype=dtype, name=name)
 File "/tensorflow/python/framework/constant_op.py", line 266, in constant
   allow_broadcast=True)
 File "/tensorflow/python/framework/constant_op.py", line 277, in _constant_impl
   return _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
 File "/tensorflow/python/framework/constant_op.py", line 302, in _constant_eager_impl
   t = convert_to_eager_tensor(value, ctx, dtype)
 File "/tensorflow/python/framework/constant_op.py", line 99, in convert_to_eager_tensor
   return ops.EagerTensor(value, ctx.device_name, dtype)
ValueError: TypeError: object of type 'RaggedTensor' has no len()


During handling of the above exception, another exception occurred:

Traceback (most recent call last):
 File "/py/absl/testing/parameterized.py", line 314, in bound_param_test
   return test_method(self, **testcase_params)
 File "/tensorflow/python/framework/test_combinations.py", line 367, in decorated
   execute_test_method()
 File "/tensorflow/python/framework/test_combinations.py", line 350, in execute_test_method
   test_method(**kwargs_to_pass)
 File "/tensorflow/python/keras/losses_test.py", line 1180, in test_ragged_tensors_3d
   loss = cce_obj(y_true, y_pred)
 File "/tensorflow/python/keras/losses.py", line 155, in __call__
   losses = call_fn(y_true, y_pred)
 File "/tensorflow/python/keras/losses.py", line 259, in call
   return ag_fn(y_true, y_pred, **self._fn_kwargs)
 File "/tensorflow/python/util/dispatch.py", line 210, in wrapper
   result = dispatch(wrapper, args, kwargs)
 File "/tensorflow/python/util/dispatch.py", line 122, in dispatch
   result = dispatcher.handle(args, kwargs)
 File "/tensorflow/python/util/dispatch.py", line 154, in handle
   return self._override_func(*args, **kwargs)
 File "/tensorflow/python/keras/losses.py", line 1762, in _ragged_tensor_sparse_categorical_crossentropy
   return _ragged_tensor_apply_loss(fn, y_true, y_pred, y_pred_extra_dim=True)
 File "/tensorflow/python/keras/losses.py", line 1291, in _ragged_tensor_apply_loss
   assertion_list = ragged_util.assert_splits_match(nested_splits_list)
 File "/tensorflow/python/ops/ragged/ragged_util.py", line 52, in assert_splits_match
   raise ValueError(error_msg)
ValueError: Inputs must have identical ragged splits

[pedro-r-marques]

@rthadur
I've updated the patch with the following change in _ragged_tensor_apply_loss:

-  if y_pred_extra_dim and y_pred.shape[-1] is None:
+  if y_pred_extra_dim:
     nested_splits_list[1] = nested_splits_list[1][:-1]

The flag y_pred_extra_dim is only set when calculating losses for sparse categorical cross entropy since we expect the shape of y_pred to have one extra dimension. This if statement had an extra check which is both unnecessary and incorrect. I would appreciate it if you could trigger the google internal CI and check that the patch passes the tests.

[pedro-r-marques]

@tomerk @rthadur Please retrigger the CI.