Home
last modified time | relevance | path

Searched refs:einsum_config (Results 1 – 6 of 6) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dmatrix.h108 absl::string_view einsum_config, int64 x_rank, int64 y_rank);
113 std::string NormalizeEinsumString(absl::string_view einsum_config);
117 xla::XlaOp x, xla::XlaOp y, absl::string_view einsum_config,
120 xla::XlaOp x, absl::string_view einsum_config,
Dmatrix.cc559 absl::string_view einsum_config, int64 x_rank, int64 y_rank) { in ParseEinsumString() argument
562 absl::StrSplit(einsum_config, ','); in ParseEinsumString()
649 std::string NormalizeEinsumString(absl::string_view einsum_config) { in NormalizeEinsumString() argument
650 if (einsum_config.find("->") != einsum_config.npos) { in NormalizeEinsumString()
653 bool has_ellipsis = einsum_config.find("...") != einsum_config.npos; in NormalizeEinsumString()
655 for (char c : einsum_config) { in NormalizeEinsumString()
660 std::string new_config(einsum_config.begin(), einsum_config.end()); in NormalizeEinsumString()
673 XlaOp Einsum(XlaOp x, XlaOp y, absl::string_view einsum_config, in Einsum() argument
677 auto new_config = NormalizeEinsumString(einsum_config); in Einsum()
685 ParseEinsumString(einsum_config, x_shape.rank(), y_shape.rank())); in Einsum()
[all …]
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Ddot_handler.cc593 [&](const WindowedEinsumConfig& einsum_config) in PartitionBaseCase()
595 CHECK(!einsum_config.windowed_at_batch_dims || in PartitionBaseCase()
596 !einsum_config.windowed_at_contracting_dims); in PartitionBaseCase()
597 const bool windowed_at_batch_dims = einsum_config.windowed_at_batch_dims; in PartitionBaseCase()
599 einsum_config.windowed_at_contracting_dims; in PartitionBaseCase()
601 einsum_config.operands_sharded_at_contracting_dims; in PartitionBaseCase()
606 einsum_config.windowed_op == WindowedEinsumOperand::LHS; in PartitionBaseCase()
/external/tensorflow/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/
Dhlo_ops.td980 StrAttr:$einsum_config
992 StrAttr:$einsum_config
/external/tensorflow/tensorflow/compiler/mlir/hlo/tests/
Dcanonicalize.mlir567 // CHECK: "mhlo.einsum"(%[[ONE]], %arg0) {einsum_config = ",ab->aa"}
568 %0 = "mhlo.unary_einsum"(%arg0) {einsum_config = "ab->aa"} : (tensor<2x3xf32>) -> tensor<2x2xf32>
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/translate/
Dexport.mlir464 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "ab,bc->ac"} : (tensor<3x4xi32>, tensor<4x5xi32>…