Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dmatrix.h119 absl::string_view einsum_config, int64_t x_rank, int64_t y_rank);
124 std::string NormalizeEinsumString(absl::string_view einsum_config);
128 xla::XlaOp x, xla::XlaOp y, absl::string_view einsum_config,
132 xla::XlaOp x, absl::string_view einsum_config,
Dmatrix.cc595 absl::string_view einsum_config, int64_t x_rank, int64_t y_rank) { in ParseEinsumString() argument
598 absl::StrSplit(einsum_config, ','); in ParseEinsumString()
685 std::string NormalizeEinsumString(absl::string_view einsum_config) { in NormalizeEinsumString() argument
686 if (einsum_config.find("->") != einsum_config.npos) { in NormalizeEinsumString()
689 bool has_ellipsis = einsum_config.find("...") != einsum_config.npos; in NormalizeEinsumString()
691 for (char c : einsum_config) { in NormalizeEinsumString()
696 std::string new_config(einsum_config.begin(), einsum_config.end()); in NormalizeEinsumString()
709 XlaOp Einsum(XlaOp x, XlaOp y, absl::string_view einsum_config, in Einsum() argument
714 auto new_config = NormalizeEinsumString(einsum_config); in Einsum()
722 ParseEinsumString(einsum_config, x_shape.rank(), y_shape.rank())); in Einsum()
[all …]
/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/tests/Dialect/mhlo/
Dhlo-legalize-einsum-to-dot-general.mlir5 …%1 = "mhlo.einsum"(%0, %arg0) {einsum_config = ",ii->i"} : (tensor<f32>, tensor<6x6xf32>) -> tenso…
17 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "bxy,bijy->bijx"} : (tensor<8x2x6xf32>, tensor<8…
35 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "ij,jk->ik"} : (tensor<?x?xf32>, tensor<?x?xf32>…
48 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "ij,j->i"} : (tensor<?x?xf32>, tensor<?xf32>) ->…
61 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "i,i->"} : (tensor<?xf32>, tensor<?xf32>) -> ten…
Dhlo-legalize-to-linalg.mlir2152 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "ijk,ikm->ijm", someattr}: (tensor<3x4x5xf32>, t…
2194 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "abc,abc->abc"} : (tensor<3x4x5xf32>, tensor<3x4…
2215 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "ae,ed->ad"}: (tensor<7x9xf32>, tensor<9x5xf32>)…
2239 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "abcdh,hg->abcdg"}: (tensor<3x4x5x6x7xf32>, tens…
2263 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "...x,xy->...y"} : (tensor<1x512x128xf32>, tenso…
2287 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "abc,cd->abd"} : (tensor<?x?x4xf32>, tensor<4x?x…
Dops.mlir4528 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "ab,bc->ac"} : (tensor<1x2xi4>, tensor<2x1xi4>) …
4536 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "ab,bc->ac"} : (tensor<1x2xi8>, tensor<2x1xi8>) …
/external/tensorflow/tensorflow/compiler/mlir/hlo/tests/
Dhlo-legalize-einsum-to-dot-general.mlir5 …%1 = "mhlo.einsum"(%0, %arg0) {einsum_config = ",ii->i"} : (tensor<f32>, tensor<6x6xf32>) -> tenso…
20 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "bxy,bijy->bijx"} : (tensor<8x2x6xf32>, tensor<8…
35 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "ij,jk->ik"} : (tensor<?x?xf32>, tensor<?x?xf32>…
50 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "ij,j->i"} : (tensor<?x?xf32>, tensor<?xf32>) ->…
65 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "i,i->"} : (tensor<?xf32>, tensor<?xf32>) -> ten…
/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/lib/Dialect/mhlo/transforms/
Dlegalize_einsum_to_dot_general.cc37 StringRef equation = einsum.einsum_config(); in matchAndRewrite()
Dlegalize_to_linalg.cc387 auto einsumConfig = op.einsum_config(); in matchAndRewrite()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Ddot_handler.cc883 [&](const WindowedEinsumConfig& einsum_config) in PartitionBaseCase()
885 CHECK(!einsum_config.windowed_at_batch_dims || in PartitionBaseCase()
886 !einsum_config.windowed_at_contracting_dims); in PartitionBaseCase()
887 const bool windowed_at_batch_dims = einsum_config.windowed_at_batch_dims; in PartitionBaseCase()
889 einsum_config.windowed_at_contracting_dims; in PartitionBaseCase()
891 einsum_config.operands_sharded_at_contracting_dims; in PartitionBaseCase()
896 einsum_config.windowed_op == WindowedEinsumOperand::LHS; in PartitionBaseCase()
/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/stablehlo/dialect/
DStablehloOps.td1714 StrAttr:$einsum_config
1726 StrAttr:$einsum_config
/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/include/mlir-hlo/Dialect/mhlo/IR/
Dhlo_ops.td1867 StrAttr:$einsum_config
1879 StrAttr:$einsum_config
/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/tests/Dialect/mhlo/canonicalize/
Dcanonicalize.mlir929 // CHECK: "mhlo.einsum"(%[[ONE]], %arg0) {einsum_config = ",ab->aa"}
930 %0 = "mhlo.unary_einsum"(%arg0) {einsum_config = "ab->aa"} : (tensor<2x3xf32>) -> tensor<2x2xf32>
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/translate/
Dexport.mlir672 …%0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "ab,bc->ac"} : (tensor<3x4xi32>, tensor<4x5xi32>…
/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/stablehlo/tests/
Dops_stablehlo.mlir4457 …%0 = "stablehlo.einsum"(%arg0, %arg1) {einsum_config = "ab,bc->ac"} : (tensor<1x2xi4>, tensor<2x1x…
4465 …%0 = "stablehlo.einsum"(%arg0, %arg1) {einsum_config = "ab,bc->ac"} : (tensor<1x2xi8>, tensor<2x1x…