Home
last modified time | relevance | path

Searched refs:iota_dimension (Results 1 – 25 of 45) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/mlir/hlo/tests/
Dlegalize-to-std.mlir117 %0 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xi32>
125 %0 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<2x4xi32>
133 %0 = "mhlo.iota"() {iota_dimension = 1 : i64} : () -> tensor<2x4xi32>
141 %0 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<2x3x4xi32>
149 %0 = "mhlo.iota"() {iota_dimension = 1 : i64} : () -> tensor<2x3x4xi32>
157 %0 = "mhlo.iota"() {iota_dimension = 2 : i64} : () -> tensor<2x3x4xi32>
165 %0 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xf32>
173 %0 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xf64>
181 %0 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xbf16>
191 %0 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xcomplex<f32>>
[all …]
Dcanonicalize.mlir483 %0 = "mhlo.dynamic_iota"(%arg0) {iota_dimension = 0 : i64} : (tensor<1xindex>) -> tensor<4xi32>
489 // CHECK: [[IOTA:%.+]] = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<5xi32>
491 %0 = "mhlo.dynamic_iota"(%arg0) {iota_dimension = 0 : i64} : (tensor<2xindex>) -> tensor<5x?xi32>
502 …// CHECK-NEXT: [[IOTA:%.+]] = "mhlo.dynamic_iota"([[CAST2]]) {iota_dimension = 0 : i64} : (tensor<…
504 %0 = "mhlo.dynamic_iota"(%arg0) {iota_dimension = 1 : i64} : (tensor<2xindex>) -> tensor<5x?xi32>
514 %0 = "mhlo.dynamic_iota"(%arg0) {iota_dimension = 0 : i64} : (tensor<2xindex>) -> tensor<1x?xi32>
523 %0 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<1xi32>
532 %0 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<1x4xi32>
542 %0 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xi32>
548 // CHECK: [[IOTA:%.+]] = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<5xi32>
[all …]
Dhlo-legalize-to-lhlo.mlir270 {iota_dimension = 0 : i64} : () -> tensor<10xi32>
271 // CHECK: "lmhlo.iota"(%{{.*}}) {iota_dimension = 0 : i64}
/external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/IR/
Dhlo_ops.cc295 auto iota_dimension = op.iota_dimension(); in Verify() local
296 if (iota_dimension >= shape.getRank() || iota_dimension < 0) in Verify()
314 auto iota_dimension = iota.iota_dimension(); in matchAndRewrite() local
317 {result_ty.getDimSize(iota_dimension)}, result_ty.getElementType()); in matchAndRewrite()
324 {iota_dimension}); in matchAndRewrite()
337 auto dimension = iota_dimension(); in fold()
363 rewriter.replaceOpWithNewOp<IotaOp>(iota, result_ty, iota.iota_dimension()); in matchAndRewrite()
380 auto iota_dimension = iota.iota_dimension(); in matchAndRewrite() local
381 auto iota_dimension_int = iota_dimension; in matchAndRewrite()
412 {iota_dimension}); in matchAndRewrite()
/external/tensorflow/tensorflow/compiler/xla/service/
Dstable_sort_expander.cc53 Cast<HloIotaInstruction>(operand)->iota_dimension() == in ExpandInstruction()
Dtopk_rewriter.cc85 iota->iota_dimension() != sort->sort_dimension()) { in Run()
Dhlo_creation_utils.h131 int64 iota_dimension);
Dhlo_sharding_util.cc1286 if (iota->iota_dimension() != index_dim) { in GetGatherBatchParallelDims()
1289 iota->iota_dimension(); in GetGatherBatchParallelDims()
1296 if (iota->iota_dimension() != index_dim) { in GetGatherBatchParallelDims()
1303 iota->iota_dimension()); in GetGatherBatchParallelDims()
Dhlo_creation_utils.cc281 int64 iota_dimension) { in MakeIotaHlo() argument
283 HloInstruction::CreateIota(shape, iota_dimension)); in MakeIotaHlo()
Dhlo_instructions.cc2918 HloIotaInstruction::HloIotaInstruction(const Shape& shape, int64 iota_dimension) in HloIotaInstruction() argument
2920 iota_dimension_(iota_dimension) {} in HloIotaInstruction()
2924 proto.add_dimensions(iota_dimension()); in ToProto()
2930 return {StrCat("iota_dimension=", iota_dimension())}; in ExtraAttributesToStringImpl()
2938 return iota_dimension() == casted_other.iota_dimension(); in IdenticalSlowPath()
2944 return absl::make_unique<HloIotaInstruction>(shape, iota_dimension()); in CloneWithNewOperandsImpl()
Dhlo_verifier.cc614 int64 iota_dimension = iota->iota_dimension(); in HandleIota() local
615 if (iota_dimension >= rank || iota_dimension < 0) { in HandleIota()
Dhlo_instructions.h1747 explicit HloIotaInstruction(const Shape& shape, int64 iota_dimension);
1749 int64 iota_dimension() const { return iota_dimension_; } in iota_dimension() function
Dalgebraic_simplifier.cc3135 dims[Cast<HloIotaInstruction>(operand)->iota_dimension()])); in HandleBroadcast()
3312 if (iota->shape().dimensions(iota->iota_dimension()) <= 1) { in HandleIota()
3824 Cast<HloIotaInstruction>(iota)->iota_dimension()); in HandleRemainder()
3852 Cast<HloIotaInstruction>(iota)->iota_dimension()); in HandleRemainder()
3933 return dim_pair.first == iota->iota_dimension() && in HandleReshape()
3938 return dim_pair.first == iota->iota_dimension() + 1; in HandleReshape()
/external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/
Dlegalize_to_standard.cc120 auto dimension = op.iota_dimension(); in matchAndRewrite()
/external/tensorflow/tensorflow/compiler/mlir/xla/ir/
Dmlir_hlo_builder.h153 XlaOp Iota(const Shape& shape, int64 iota_dimension) override;
Dmlir_hlo_builder.cc201 XlaOp MlirHloBuilder::Iota(const Shape& shape, int64 iota_dimension) { in Iota() argument
208 builder_.getIntegerAttr(builder_.getI64Type(), iota_dimension)); in Iota()
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/translate/
Dimport.hlotxt460 // CHECK-NEXT: "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xf32>
461 ROOT %iota.0 = f32[4] iota(), iota_dimension=0
466 // CHECK-NEXT: "mhlo.iota"() {iota_dimension = 1 : i64} : () -> tensor<4x5xf32>
467 ROOT %iota.0 = f32[4, 5] iota(), iota_dimension=1
Dexport.mlir542 iota_dimension = 1 : i64
548 // CHECK: ROOT %[[RESULT:.*]] = f32[1,10] iota(), iota_dimension=1
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
Dlegalize_hlo.cc719 const uint64_t dimension = iota_op.iota_dimension(); in matchAndRewrite()
744 reshape_shape[iota_op.iota_dimension()] = type.getShape()[dimension]; in matchAndRewrite()
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/
Dlegalize-tf.mlir546 // CHECK-DAG: %[[IOTA0:.*]] = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<12x12xi32>
547 // CHECK-DAG: %[[IOTA1:.*]] = "mhlo.iota"() {iota_dimension = 1 : i64} : () -> tensor<12x12xi32>
570 // CHECK-DAG: %[[V2:.*]] = "mhlo.iota"() {iota_dimension = 1 : i64} : () -> tensor<1x22x128xi32>
571 // CHECK-DAG: %[[V3:.*]] = "mhlo.iota"() {iota_dimension = 2 : i64} : () -> tensor<1x22x128xi32>
1257 // CHECK-DAG: %[[X:.*]] = "mhlo.iota"() {iota_dimension = 1 : i64} : () -> tensor<64x64xi64>
1258 // CHECK-DAG: %[[Y:.*]] = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<64x64xi64>
1277 // CHECK-DAG: %[[X:.*]] = "mhlo.iota"() {iota_dimension = 1 : i64} : () -> tensor<24x48xi64>
1278 // CHECK-DAG: %[[Y:.*]] = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<24x48xi64>
1439 // CHECK: %[[IOTA:.*]] = "mhlo.iota"() {iota_dimension = 1 : i64} : () -> tensor<3x5xi32>
3371 // CHECK: %[[INDEX:.*]] = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<3x7xi32>
[all …]
/external/tensorflow/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/
Dhlo_ops.td75 let arguments = (ins I64Attr:$iota_dimension);
95 let arguments = (ins HLO_DimensionTensor:$output_shape, I64Attr:$iota_dimension);
Dlhlo_ops.td65 let arguments = (ins I64Attr:$iota_dimension,
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/hlo_to_lhlo_with_xla/
Dops.mlir782 // CHECK: "lmhlo.iota"(%{{.*}}) {iota_dimension = 1 : i64} : (memref<1x10xf32>) -> ()
785 iota_dimension = 1 : i64
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_builder.h782 virtual XlaOp Iota(const Shape& shape, int64 iota_dimension);
1355 int64 iota_dimension);
2342 XlaOp Iota(XlaBuilder* builder, const Shape& shape, int64 iota_dimension);
/external/tensorflow/tensorflow/compiler/mlir/xla/
Dhlo_function_importer.cc255 Cast<HloIotaInstruction>(instruction)->iota_dimension())) in ImportInstructionImpl()

12