| /external/tensorflow/tensorflow/core/kernels/ |
| D | dynamic_partition_op_gpu.cu.cc | 209 Tensor* partitions_out, Tensor* indices_out, in AllocateTempSpace() argument
219 c, c->allocate_temp(DT_INT32, TensorShape({N}), indices_out), done); in AllocateTempSpace()
278 Tensor indices_out; in ComputeAsync() local
282 this->CountAndSortParts(c, &partitions, &partition_count, &indices_out, in ComputeAsync()
315 auto wrapped_callback = [this, c, &data, &partitions, indices_out, in ComputeAsync()
325 this->GatherSlices(c, &data, &indices_out, N, slice_size, outputs); in ComputeAsync()
337 Tensor* indices_out, DoneCallback done) { in RadixSort() argument
348 int32* indices_out_ptr = indices_out->flat<int32>().data(); in RadixSort()
370 Tensor* partition_count, Tensor* indices_out, in CountAndSortParts() argument
380 this->AllocateTempSpace(c, N, &indices_in, &partitions_out, indices_out, in CountAndSortParts()
[all …]
|
| D | sparse_cross_op.cc | 143 Tensor* indices_out, Tensor* values_out) in OutputUpdater() argument
145 indices_out_(indices_out), in OutputUpdater()
320 Tensor* indices_out; in Compute() local
325 CreateOutputTensors(columns, batch_size, context, &indices_out, &values_out, in Compute()
329 output_start_indices, indices_out, values_out); in Compute()
506 int64 batch_size, OpKernelContext* context, Tensor** indices_out, in CreateOutputTensors() argument
523 0, TensorShape({cross_count_total, 2}), indices_out)); in CreateOutputTensors()
|
| D | topk_op.cc | 81 Tensor* indices_out = nullptr; in Compute() local
83 context->allocate_output(1, output_shape, &indices_out)); in Compute()
89 auto indices = indices_out->flat_inner_dims<int32>(); in Compute()
|
| /external/tensorflow/tensorflow/python/client/ |
| D | session_test.py | 626 indices_out, values_out, shape_out = sp_out
627 self.assertAllEqual(indices_out, indices)
636 indices_out, values_out, shape_out = s.run(sp)
637 self.assertAllEqual(indices_out, indices)
641 (indices_out, values_out, shape_out), = s.run([sp])
642 self.assertAllEqual(indices_out, indices)
651 indices_out, values_out, shape_out = s.run({'sp': sp})['sp']
652 self.assertAllEqual(indices_out, indices)
656 (indices_out, values_out, shape_out), = s.run({'sp': [sp]})['sp']
657 self.assertAllEqual(indices_out, indices)
[all …]
|
| /external/tensorflow/tensorflow/core/kernels/sparse/ |
| D | kernels_gpu.cu.cc | 184 int64* indices_out, int size) { in COOMatrixToSparseTensorKernel2D() argument
186 indices_out[i * 2] = static_cast<int64>(ldg(coo_rows + i)); in COOMatrixToSparseTensorKernel2D()
187 indices_out[i * 2 + 1] = static_cast<int64>(ldg(coo_cols + i)); in COOMatrixToSparseTensorKernel2D()
207 const int* coo_rows, const int* coo_cols, int64* indices_out, in COOMatrixToSparseTensorKernel3D() argument
223 indices_out[i * 3] = static_cast<int64>(b); in COOMatrixToSparseTensorKernel3D()
224 indices_out[i * 3 + 1] = static_cast<int64>(ldg(coo_rows + i)); in COOMatrixToSparseTensorKernel3D()
225 indices_out[i * 3 + 2] = static_cast<int64>(ldg(coo_cols + i)); in COOMatrixToSparseTensorKernel3D()
|
| /external/tensorflow/tensorflow/python/ops/ |
| D | while_v2_indexed_slices_rewriter.py | 172 indices_out = grad_output_slices.indices
190 indices = constant_op.constant([], indices_out.dtype, name="indices_init")
|
| D | sparse_ops.py | 643 indices_out, values_out, shape_out = gen_sparse_ops.sparse_cross(
655 return sparse_tensor.SparseTensor(indices_out, values_out, shape_out)
|