diff options
author | 2017-09-27 13:48:03 -0700 | |
---|---|---|
committer | 2017-09-27 13:53:53 -0700 | |
commit | 2a5fb08bf2885cba29065d7269c5f6a32614b89a (patch) | |
tree | f8bf4a817bdce4daecdbb22a897f7587985c18f7 /tensorflow/core | |
parent | 301b14c240fe99249dc2225132a7ebe5cbecbdc4 (diff) |
SymbolicGradients: create the underlying runtime with the correct step container.
This fixes a bug where calling tf.gradients of a tf.while_loop inside a Defun
would hard crash the program.
Also added some safety checks inside StackOps to avoid the hard crash if
something like this happens again.
PiperOrigin-RevId: 170246274
Diffstat (limited to 'tensorflow/core')
-rw-r--r-- | tensorflow/core/kernels/function_ops.cc | 1 | ||||
-rw-r--r-- | tensorflow/core/kernels/stack_ops.cc | 11 |
2 files changed, 10 insertions, 2 deletions
diff --git a/tensorflow/core/kernels/function_ops.cc b/tensorflow/core/kernels/function_ops.cc index a7206f6258..584d41dfe0 100644 --- a/tensorflow/core/kernels/function_ops.cc +++ b/tensorflow/core/kernels/function_ops.cc @@ -241,6 +241,7 @@ class SymbolicGradientOp : public AsyncOpKernel { opts.cancellation_manager = ctx->cancellation_manager(); opts.runner = ctx->runner(); opts.stats_collector = ctx->stats_collector(); + opts.step_container = ctx->step_container(); std::vector<Tensor> args; args.reserve(ctx->num_inputs()); for (int i = 0; i < ctx->num_inputs(); ++i) { diff --git a/tensorflow/core/kernels/stack_ops.cc b/tensorflow/core/kernels/stack_ops.cc index a474e75d6a..affe81a555 100644 --- a/tensorflow/core/kernels/stack_ops.cc +++ b/tensorflow/core/kernels/stack_ops.cc @@ -150,7 +150,11 @@ Status GetStack(OpKernelContext* ctx, Stack** stack) { if (rm == nullptr) { return errors::Internal("No resource manager."); } - TF_RETURN_IF_ERROR(rm->Lookup(ctx->step_container()->name(), key, stack)); + auto* step_container = ctx->step_container(); + if (step_container == nullptr) { + return errors::Internal("No step container."); + } + TF_RETURN_IF_ERROR(rm->Lookup(step_container->name(), key, stack)); return Status::OK(); } @@ -191,7 +195,10 @@ class StackOp : public OpKernel { OP_REQUIRES(ctx, rm != nullptr, errors::Internal("No resource manager.")); string key = strings::StrCat(kContainer, stack_name); Stack* stack = new Stack(elem_type_, stack_name, size); - OP_REQUIRES_OK(ctx, rm->Create(ctx->step_container()->name(), key, stack)); + auto* step_container = ctx->step_container(); + OP_REQUIRES(ctx, step_container != nullptr, + errors::Internal("No step container.")); + OP_REQUIRES_OK(ctx, rm->Create(step_container->name(), key, stack)); if (IsRefType(ctx->expected_output_dtype(0))) { // Create the stack handle. AllocatorAttributes alloc_attr; |