Searched defs:batch_indices (Results 1 – 9 of 9) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/core/kernels/ |
| H A D | gather_nd_op_gpu.cu.cc | 33 const Eigen::array<int64, IXDIM> batch_indices, const int64 indices_size, in GatherSliceOpKernel()
78 Eigen::array<int64, IXDIM> batch_indices; in operator ()() local
|
| H A D | ctc_loss_op.cc | 162 const int64_t batch_indices = g.group()[0]; in Compute() local
|
| H A D | topk_op_gpu.h | 362 auto batch_indices = indices + offset; in TopKKernel() local
|
| H A D | batch_kernels.cc | 674 auto batch_indices = in Compute() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/util/ |
| H A D | bcast.h | 107 const std::vector<int64_t>& batch_indices(int i) const { in batch_indices() function
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/ |
| H A D | TensorConversions.cpp | 699 auto batch_indices = at::arange(0, n_batch, options).repeat_interleave(nnz_per_batch); in sparse_compressed_to_dense() local
1560 auto batch_indices = indices.narrow(0, 0, batch_ndim); in convert_indices_from_csr_to_coo_cpu() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| H A D | resampler_ops.cc | 130 auto batch_indices = in ConcatenateIota() local
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/sparse/cuda/ |
| H A D | SparseCsrTensorMath.cu | 110 auto batch_indices = indices.narrow(0, 0, batch_ndim); in convert_indices_from_csr_to_coo_cuda() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/cc/gradients/ |
| H A D | array_grad.cc | 633 Output batch_indices = indices; in GetBatchIndices() local
|