Searched defs:batch_idx (Results 1 – 15 of 15) sorted by relevance
91 for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) { in operator ()() local
307 [&](int64 batch_idx, int64 row_begin, int64 row_end) { in SparseDenseMatMulWithoutTransposedLHS()378 [&](int64 batch_idx, int64 row_begin, int64 row_end) { in SparseDenseMatMulWithTransposedLHS()433 for (int64 batch_idx = batch_begin; batch_idx <= batch_end_inclusive; in HandleBatchAndRowRange() local
187 for (int64 batch_idx = batch_begin; batch_idx < batch_end; in Compute() local229 for (int64 batch_idx = batch_begin; batch_idx < batch_end; in Compute() local
96 for (int64 batch_idx = batch_begin; batch_idx < batch_end; ++batch_idx) { in Compute() local
100 for (int64 batch_idx = batch_begin; batch_idx < batch_end; ++batch_idx) { in Compute() local
76 xla::XlaOp batch_idx = xla::Iota( in Compile() local
58 for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) { in SliceInput() local211 for (int batch_idx = 0; batch_idx < bcast.output_batch_size(); ++batch_idx) { in Run() local
118 uint32_t batch_idx = set_bit_count_until_i++ % kBatchSize; in ReadSetBitBatch() local
105 for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) { in sliceInput() local167 for (int batch_idx = 0; batch_idx < bcast.output_batch_size(); ++batch_idx) { in createMatMulOps() local
61 SliceIndex batch_idx = static_cast<SliceIndex>(start / indices_size); in HandleCopies() local
63 SliceIndex batch_idx = static_cast<SliceIndex>( in HandleCopiesBatched() local
199 int64 batch_idx = output_idx / samples_per_batch; in operator ()() local
736 for (int64 batch_idx = 0, dest_idx = 0; batch_idx < batch_size; in AddBatchOffsets() local
190 int depth_idx, batch_idx, spatial_idx; in GetDimIndices() local
68 int batch_idx; member