Note that @benbarsdell will soon open a PR to make tf.math.unsorted_segment_sum operate deterministically on a GPU. This will remove the d9m-unimplemented exception-throwing for this op when running on a GPU (and when op-determinism is expected). This current PR cannot be merged until that exception-throwing functionality has been removed from the top-of-tree.

Perhaps this should be implemented within the C++ op instead of on the Python level, which would solve the issue of raising an exception on the CPU. The core logic of the op, which are the following two lines, can be done in the C++ SparseSoftmaxCrossEntropyWithLogits op when determinism is enabled:

    log_probs = log_softmax_v2(logits)
    cost = math_ops.negative(array_ops.gather(log_probs, labels, batch_dims=1))

I think the current partially-implemented approach of checking math_ops.add(1, 2).device cannot be made to work in all cases, such as in graph mode when a GraphDef is saved then run on a separate machine. Alternatively, we could just always return NaNs even on the CPU, which is what XLA does despite it violating the documentation, but I would prefer raising exceptions on the CPU.

Implementing in C++ has a few more advantages. We don't have to deal with the awkwardness of a custom gradient. It's also probably faster as we don't have to run a lot of ops (like tf.where, tf.logical_or, etc). I also think it's slightly cleaner for the same op to be run regardless of whether determinism is enabled (we already run different ops for the non-sparse version but in that case the logic is a lot simpler).

The main disadvantage that I can think of is implementing the op in C++ will take more time. Unfortunately most of the non-test code in this PR would not be used.

@duncanriach what do you think? To implement, you could call the GatherFunctor and the SoftmaxFunctor from the SparseSoftmaxCrossEntropyWithLogits op.

Also /CC @sanjoy

@reedwm, the NaN injection (in both the forward and backward directions) when running on GPU, is currently provided by the Eigen-based implementation (in SparseXentLossGenerator and SparseXentGradGenerator in sparse_xent_op.h). Calling GatherFunctor and SoftmaxFunctor in SparseSoftmaxXentWithLogitsOp (implementing the math as in the Python-level solution) would not properly implement the NaN injection (or even the backprop output from the op). It would, however, solve the problem of throwing the exception when running on CPU. Maybe I'm misunderstanding something about your proposal; let me know if you think I've missed something.

I think that two reasonable (relatively temporary) compromises would be:

1. When op-determinism is enabled, inject NaNs when running on CPU as well as GPU (already basically implemented).
2. When op-determinism is enabled, neither inject NaNs nor throw an exception on either CPU or GPU (already basically implemented).

The second option would improve performance and is similar to what I did for tf.nn.bias_add, the current deterministic implementation of which skips some error checks (see here).

I suspect that for these ops for which we provide interim, deterministic Python-level solutions, it may be possible to later implement higher-performance, CUDA-level solutions that also include complete, high-performance error checks.

Some responses to your overall review comment, @reedwm:

We could get around this by adding a parameter to GatherFunctor to inject NaN instead of zeros.

True. However, this does not address the backprop path that would need to be implemented under SparseSoftmaxXentWithLogitsOp as well, including the injection of NaN into the parts of the gradient associated with illegal label values.

I prefer (2) since the implementation is simpler, there is precedent in tf.nn.bias_add, and we can eventually solve this by moving the determinism codepath of this op to C++.

Right. And the details of how that's implemented at the C++ level can be worked out later. I plan to include one or more TODOs.

On the CPU, an exception will still be raised with (2) because tf.gather raises an exception on the CPU.

This doesn't match with my memory of what I saw in the testing (no exception on CPU), and I cannot find this exception-raising functionality in the code. I do see that the documentation claims that there will be an exception raised on the CPU. I'll review this while revising the tests.

On the GPU, zeros will be returned instead of NaN if the index is out of bounds due to tf.gather returning zeros.

Yes, and I seem to remember that zeros are also returned by the backprop of the deterministic solution (rather than NaNs), which I think makes sense mathematically: when an invalid logit is the "correct" one, then changes in the actual logits will have no effect on the loss.

The following code will be removed from this pull request (from python/ops/nn_ops.py) by the next commit. It enables the python-level, GPU-deterministic implementation of tf.nn.spase_softmax_cross_entropy_with_logits to inject NaNs into the batch entries for which the labels value is invalid (in both the forward and backward direction). I'm placing this code here for potential future reference.

@custom_gradient.custom_gradient
def _sparse_softmax_cross_entropy_gradient_nan_injector(logits, labels):
  logits_shape = array_ops.shape(logits)
  class_count = math_ops.cast(logits_shape[-1], labels.dtype)
  nan_tensor = constant_op.constant(float('Nan'), dtype=logits.dtype)
  def grad(upstream_gradient):
    logits_all_nans = array_ops.broadcast_to(nan_tensor, logits_shape)
    labels_on_logits = array_ops.transpose(
        labels * array_ops.transpose(
            [array_ops.ones(logits_shape[1], dtype=labels.dtype)]))
    logits_grad = array_ops.where(
        math_ops.logical_or(
            math_ops.less(labels_on_logits, 0),
            math_ops.greater_equal(labels_on_logits, class_count)),
        logits_all_nans, upstream_gradient)
    labels_grad = None
    return [logits_grad, labels_grad]
  return logits, grad


def _sparse_softmax_cross_entropy_with_rank_2_logits(logits, labels, name):
  if _tf_deterministic_ops():
    # Handle invalid labels before running the op
    _sparse_softmax_cross_entropy_check_for_valid_labels_on_cpu(logits, labels)
    logits = _sparse_softmax_cross_entropy_gradient_nan_injector(logits, labels)

    # The actual op functionality
    log_probs = log_softmax_v2(logits)
    cost = math_ops.negative(array_ops.gather(log_probs, labels, batch_dims=1))

    # Force the output to be NaN when the corresponding label is invalid
    nan_tensor = constant_op.constant(float('Nan'), dtype=logits.dtype)
    cost_all_nans = array_ops.broadcast_to(nan_tensor, array_ops.shape(cost))
    class_count = math_ops.cast(array_ops.shape(logits)[-1], labels.dtype)
    cost = array_ops.where(
        math_ops.logical_or(
          math_ops.less(labels, 0),
          math_ops.greater_equal(labels, class_count)),
        cost_all_nans, cost)
  else:
    # The second output tensor contains the gradients. We use it in
    # _CrossEntropyGrad() in nn_grad but not here.
    cost, _ = gen_nn_ops.sparse_softmax_cross_entropy_with_logits(
        logits, labels, name=name)
  return cost

Here is the difference in the gradient generated by testInvalidLabelGPU when the selective gradient gating after tf.gather is removed:

Expected gradient:

[[ 0.    ,  0.    ,  0.    ,  0.    ],
 [ 0.25  ,  0.25  ,  0.25  , -0.75  ],
 [-0.968 ,  0.087 ,  0.237 ,  0.6439],
 [ 0.    ,  0.    ,  0.    ,  0.    ]])

Actual gradient:

[[ 0.      ,  0.      ,  0.      ,  0.      ],
 [-0.5     ,  0.5     ,  0.5     , -0.5     ],
 [-0.935883,  0.174289,  0.473766,  0.287828],
 [ 0.      ,  0.      ,  0.      ,  0.      ]]

I'll mark this PR as approved once the unsorted_segment_sum is deterministic. The PR looks good to me and I have no further suggested changes. I don't want to approve now since someone may try to merge the PR and get confused by the "DO NOT SUBMIT".

The difference in gradients when the gating is disabled makes me suspect the tf.gather gradient or some other op might have a bug, as the NaN injection shouldn't change the gradients (unless I did the math wrong). Are you planning at looking at this? If not, I'll take a look myself.

I'll mark this PR as approved once the unsorted_segment_sum is deterministic.

Thanks. Sounds good.

The difference in gradients when the gating is disabled makes me suspect the tf.gather gradient or some other op might have a bug, as the NaN injection shouldn't change the gradients (unless I did the math wrong).

I suspect that this is revealing a bug in tf.math.unsorted_segment_sum (used with the backprop of tf.gather): the "illegal" segment ID is probably not getting stopped before it gets to tf.math.unsorted_segment_sum and then address aliasing is leading to gradients from the batch indices associated with the illegal labels spilling over, and getting added into, the batch indices associated with legal labels.

However, when the gradients associated with legal labels are blocked by injecting either zero or NaN in the forward direction (as this PR currently does), the associated gradient becomes zero and the gradients associated with legal labels come through correctly.

So, this is probably a bug associated with tf.gather (on GPU) such that while illegal values in its indices parameter lead to zeros being output in the forward direction, the associated back-propped gradients (channeled via tf.math.unsorted_segment_sum) are not being dropped appropriately.

Are you planning at looking at this? If not, I'll take a look myself.

I intend to look more deeply into this at some point (I added a TODO), but not right now. I planned to, at least, wait until after the new GPU implementation of tf.math.unsorted_segment_sum was in the top-of-tree, just in case that happened to resolve it. However, you're welcome to investigate sooner.

@duncanriach Can you please resolve conflicts? Thanks!

@gbaned, force-pushed after rebase to resolve conflicts.

@duncanriach Can you please resolve conflicts? Thanks!

@gbaned: rebased, conflicts resolved, and force-pushed.

I'm going to rebase to get the other changes in my local tree and run a test before force pushing, which will probably happen tomorrow.

@reedwm, I rebased and force-pushed, but you only need to review the most recent commit.

Seems auto-merge is not happening but the changes are merged into master now, so we can close this. Thank you for the PR.

Seems like this test is failing on Windows (we don't have log because the output is too big).

I think this is because unsorted_segment_sum is not deterministic on Windows and so will raise an error when determinism is enabled. I'll disable the test on Windows for now.