| 402 | } |
| 403 | |
| 404 | Tensor OpKernelContext::mutable_input(int index, bool lock_held) { |
| 405 | CHECK_GE(index, 0); |
| 406 | CHECK_LT(index, num_inputs()); |
| 407 | CHECK(input_is_ref(index)); |
| 408 | // return a copy of the Ref acquired while holding the mutex |
| 409 | if (lock_held) { |
| 410 | Tensor& tensor = *((*params_->inputs)[index].tensor); |
| 411 | record_tensor_reference(tensor); |
| 412 | return tensor; |
| 413 | } else { |
| 414 | tf_shared_lock l(*input_ref_mutex(index)); |
| 415 | Tensor& tensor = *((*params_->inputs)[index].tensor); |
| 416 | record_tensor_reference(tensor); |
| 417 | return tensor; |
| 418 | } |
| 419 | } |
| 420 | |
| 421 | void OpKernelContext::replace_ref_input(int index, const Tensor& tensor, |
| 422 | bool lock_held) { |