| 892 | string DebugString() const final { return "UnbatchResource"; } |
| 893 | |
| 894 | Status Compute(OpKernelContext* context, AsyncOpKernel::DoneCallback done) { |
| 895 | const Tensor& data_t = context->input(0); |
| 896 | const Tensor& batch_index_t = context->input(1); |
| 897 | |
| 898 | if (batch_index_t.shape().dim_size(0) > data_t.shape().dim_size(0)) { |
| 899 | return errors::InvalidArgument( |
| 900 | "Wrong shape for index tensor. Expected 0th dimension size to be no " |
| 901 | "greater than ", |
| 902 | data_t.shape().dim_size(0), |
| 903 | "; Got: ", batch_index_t.shape().dim_size(0), "."); |
| 904 | } |
| 905 | if (batch_index_t.shape().dim_size(1) != 3) { |
| 906 | return errors::InvalidArgument( |
| 907 | "Wrong shape for index tensor. Expected 1st dimension size to be 3 ; " |
| 908 | "Got: ", |
| 909 | batch_index_t.shape().dim_size(1), "."); |
| 910 | } |
| 911 | |
| 912 | const int64 batch_key = context->input(2).scalar<int64>()(); |
| 913 | const bool nonempty_input = batch_index_t.dim_size(0) > 0; |
| 914 | |
| 915 | // If we have a non-empty tensor, slice it up. |
| 916 | // (It is important to do this outside of the critical section below.) |
| 917 | // The following variables are populated iff 'nonempty_input==true'. |
| 918 | std::vector<int64> sizes; |
| 919 | std::vector<int64> batch_keys; |
| 920 | std::vector<Tensor> split_inputs; |
| 921 | if (nonempty_input) { |
| 922 | auto batch_indices = |
| 923 | batch_index_t.shaped<int64, 2>({batch_index_t.dim_size(0), 3}); |
| 924 | for (int i = 0; i < batch_index_t.dim_size(0); ++i) { |
| 925 | sizes.push_back(batch_indices(i, 2) - batch_indices(i, 1)); |
| 926 | batch_keys.push_back(batch_indices(i, 0)); |
| 927 | } |
| 928 | |
| 929 | const DataType type = data_t.dtype(); |
| 930 | switch (type) { |
| 931 | #define CASE(type) \ |
| 932 | case DataTypeToEnum<type>::value: \ |
| 933 | TF_RETURN_IF_ERROR(Split<type>(context, data_t, sizes, &split_inputs)); \ |
| 934 | break; |
| 935 | TF_CALL_ALL_TYPES(CASE); |
| 936 | #undef CASE |
| 937 | default: |
| 938 | return errors::InvalidArgument("Unsupported data type: ", type); |
| 939 | } |
| 940 | } |
| 941 | |
| 942 | // Critical section. |
| 943 | std::vector<AsyncOpKernel::DoneCallback> done_callbacks_to_call; |
| 944 | Status status = [&]() -> Status { |
| 945 | mutex_lock ml(mu_); |
| 946 | |
| 947 | // Check to see whether the tensor we want is already ready. |
| 948 | auto tensor_it = waiting_tensors_.find(batch_key); |
| 949 | if (tensor_it != waiting_tensors_.end()) { |
| 950 | context->set_output(0, tensor_it->second.tensor); |
| 951 | waiting_tensors_.erase(tensor_it); |