Searched refs:batch_dim_ids (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | cusolver_rewriter.cc | 62 std::vector<int64> batch_dim_ids(batch_dims.size()); in CreateCholesky() local
63 absl::c_iota(batch_dim_ids, 0); in CreateCholesky()
113 /*broadcast_dimensions=*/batch_dim_ids)); in CreateCholesky()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | qr_expander.cc | 87 std::vector<int64> batch_dim_ids(batch_dims.size()); in House() local
88 std::iota(batch_dim_ids.begin(), batch_dim_ids.end(), 0); in House()
147 *v = e_k + Div(x_after_k, divisor, /*broadcast_dimensions=*/batch_dim_ids); in House()
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | legalize_tf.cc | 5711 llvm::SmallVector<int64_t, 4> batch_dim_ids(batch_dims.size()); in House() local
5712 std::iota(batch_dim_ids.begin(), batch_dim_ids.end(), 0); in House()
5786 GetI64ElementsAttr(batch_dim_ids, builder), in House()
|