Searched refs:batch_dim_indices (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | qr_expander.cc | 195 std::vector<int64> batch_dim_indices(num_batch_dims); in QrBlock() local
196 std::iota(batch_dim_indices.begin(), batch_dim_indices.end(), 0); in QrBlock()
226 /*broadcast_dimensions=*/batch_dim_indices); in QrBlock()
238 /*broadcast_dimensions=*/batch_dim_indices); in QrBlock()
241 /*broadcast_dimensions=*/ConcatVectors(batch_dim_indices, {minor_dim})); in QrBlock()
299 std::vector<int64> batch_dim_indices(batch_dims.size()); in CompactWYRepresentation() local
300 std::iota(batch_dim_indices.begin(), batch_dim_indices.end(), 0); in CompactWYRepresentation()
321 ConcatVectors(batch_dim_indices, {n_index})); in CompactWYRepresentation()
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | legalize_tf.cc | 5827 llvm::SmallVector<int64_t, 4> batch_dim_indices(batch_dims.size()); in QRBlock() local
5828 std::iota(batch_dim_indices.begin(), batch_dim_indices.end(), 0); in QRBlock()
5862 loc, tau, vva, GetI64ElementsAttr(batch_dim_indices, builder), in QRBlock()
5894 loc, beta, mask, GetI64ElementsAttr(batch_dim_indices, builder), in QRBlock()
5997 llvm::SmallVector<int64_t, 4> batch_dim_indices(batch_dims.size()); in ComputeWYRepresentation() local
5998 std::iota(batch_dim_indices.begin(), batch_dim_indices.end(), 0); in ComputeWYRepresentation()
6046 auto beta_broadcast_dims = llvm::to_vector<4>(batch_dim_indices); in ComputeWYRepresentation()
6066 auto beta_broadcast_dims = llvm::to_vector<4>(batch_dim_indices); in ComputeWYRepresentation()
|