| 562 | } |
| 563 | |
| 564 | void Compute(OpKernelContext* c) override { |
| 565 | const TensorShape& tl_a_shape = c->input(0).shape(); |
| 566 | const TensorShape& tl_b_shape = c->input(1).shape(); |
| 567 | OP_REQUIRES( |
| 568 | c, tl_a_shape == tl_b_shape, |
| 569 | errors::InvalidArgument("Incompatible input TensorList tensor shapes: ", |
| 570 | tl_a_shape.DebugString(), " vs. ", |
| 571 | tl_b_shape.DebugString())); |
| 572 | AllocatorAttributes attr; |
| 573 | std::unique_ptr<Tensor> tl_alias = c->forward_input( |
| 574 | 0 /*input_index*/, 0 /*output_index*/, DT_VARIANT, tl_a_shape, |
| 575 | DEVICE_MEMORY /* input is always on DEVICE_MEMORY */, attr); |
| 576 | |
| 577 | // tl_a may be aliased by tl_alias. |
| 578 | const Tensor& tl_a = c->input(0); |
| 579 | const Tensor& tl_b = c->input(1); |
| 580 | |
| 581 | Tensor* output = nullptr; |
| 582 | bool ok_to_alias = tl_alias != nullptr; |
| 583 | if (tl_alias && tl_alias->dtype() == DT_VARIANT && |
| 584 | tl_alias->NumElements() > 0) { |
| 585 | auto tl_a_t = tl_alias->flat<Variant>(); |
| 586 | for (int64 i = 0; i < tl_alias->NumElements(); ++i) { |
| 587 | TensorList* aliased = tl_a_t(i).get<TensorList>(); |
| 588 | if (aliased == nullptr || !aliased->RefCountIsOne()) { |
| 589 | ok_to_alias = false; |
| 590 | break; |
| 591 | } |
| 592 | } |
| 593 | if (ok_to_alias) { |
| 594 | c->set_output(0, *tl_alias); |
| 595 | output = tl_alias.get(); |
| 596 | } |
| 597 | } |
| 598 | if (!ok_to_alias) { |
| 599 | // Couldn't alias the entire Tensor. We'll be conservative and not try |
| 600 | // to alias individual batch entries. |
| 601 | attr.set_on_host(true); |
| 602 | OP_REQUIRES_OK(c, c->allocate_output(0, tl_a_shape, &output, attr)); |
| 603 | } |
| 604 | |
| 605 | auto output_t = output->flat<Variant>(); |
| 606 | auto tl_a_t = tl_a.flat<Variant>(); |
| 607 | auto tl_b_t = tl_b.flat<Variant>(); |
| 608 | |
| 609 | for (int64 i = 0; i < tl_a.NumElements(); ++i) { |
| 610 | const TensorList* l_a = tl_a_t(i).get<TensorList>(); |
| 611 | const TensorList* l_b = tl_b_t(i).get<TensorList>(); |
| 612 | OP_REQUIRES( |
| 613 | c, l_a != nullptr, |
| 614 | errors::InvalidArgument("input_a is not a TensorList at index ", i, |
| 615 | ". Saw: '", tl_a_t(i).DebugString(), "'")); |
| 616 | OP_REQUIRES( |
| 617 | c, l_b != nullptr, |
| 618 | errors::InvalidArgument("input_b is not a TensorList at index ", i, |
| 619 | ". Saw: '", tl_b_t(i).DebugString(), "'")); |
| 620 | OP_REQUIRES(c, l_a->element_dtype == element_dtype_, |
| 621 | errors::InvalidArgument( |
nothing calls this directly
no test coverage detected