Searched refs:batch_indices (Results 1 – 12 of 12) sorted by relevance
| /external/tensorflow/tensorflow/core/kernels/ |
| D | gather_nd_op_gpu.cu.cc | 33 const Eigen::array<int64, IXDIM> batch_indices, const int64 indices_size, in GatherSliceOpKernel() argument
45 out_of_bounds |= !FastBoundsCheck(index_j, batch_indices[j]); in GatherSliceOpKernel()
78 Eigen::array<int64, IXDIM> batch_indices; in operator ()() local
81 batch_indices[size_t(IXDIM - 1)] = Tparams.dimension(IXDIM - 1); in operator ()()
84 batch_indices[i - 1] = Tparams.dimension(i - 1); in operator ()()
92 Tout.data(), batch_strides, batch_indices, in operator ()()
|
| D | ctc_loss_op.cc | 153 const int64 batch_indices = g.group()[0]; in Compute() local
154 OP_REQUIRES(ctx, FastBoundsCheck(batch_indices, batch_size), in Compute()
157 " but saw: ", batch_indices)); in Compute()
160 std::vector<int>* b_values = &labels_t[batch_indices]; in Compute()
|
| D | batch_kernels.cc | 588 auto batch_indices = in Compute() local
591 sizes.push_back(batch_indices(i, 2) - batch_indices(i, 1)); in Compute()
592 batch_keys.push_back(batch_indices(i, 0)); in Compute()
|
| D | parameterized_truncated_normal_op.cc | 360 const auto& means_batch_indices = bcast.batch_indices(0); in operator ()()
361 const auto& stddevs_batch_indices = bcast.batch_indices(1); in operator ()()
362 const auto& minvals_batch_indices = bcast.batch_indices(2); in operator ()()
363 const auto& maxvals_batch_indices = bcast.batch_indices(3); in operator ()()
|
| D | topk_op_gpu.h | 370 auto batch_indices = indices + offset;
377 batch_indices);
|
| /external/tensorflow/tensorflow/python/kernel_tests/ |
| D | lu_op_test.py | 111 batch_indices = math_ops.cast(
117 array_ops.stack([batch_indices, inv_perm_reshaped], axis=-1))
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | resampler_ops.cc | 127 auto batch_indices = in ConcatenateIota() local
131 return xla::ConcatInDim(b, {batch_indices, indices}, dimensions.size() - 1); in ConcatenateIota()
|
| /external/tensorflow/tensorflow/core/util/ |
| D | bcast.h | 107 const std::vector<int64>& batch_indices(int i) const { in batch_indices() function
|
| /external/tensorflow/tensorflow/python/ops/ |
| D | array_grad.py | 587 batch_indices = indices
601 batch_indices += array_ops.reshape(dim_indices, dim_shape)
603 return batch_indices
|
| D | array_ops.py | 5160 batch_indices = indices
5174 batch_indices += reshape(dim_indices, dim_shape)
5176 flat_indices = reshape(batch_indices, [-1])
|
| /external/tensorflow/tensorflow/python/ops/linalg/ |
| D | linalg_impl.py | 1140 batch_indices = array_ops.broadcast_to(
1142 x = array_ops.gather_nd(x, array_ops.stack([batch_indices, perm],
|
| /external/tensorflow/tensorflow/python/kernel_tests/linalg/sparse/ |
| D | csr_sparse_matrix_ops_test.py | 78 batch_indices = array_ops.broadcast_to(
83 array_ops.stack([batch_indices, permutation_indices_inv], axis=-1))
|