Home
last modified time | relevance | path

Searched refs:partitioned_output (Results 1 – 5 of 5) sorted by relevance

/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/
Dxla_sharding_util.cc553 mlir::TF::TPUPartitionedOutputOp partitioned_output; in RemapOutputsFromLogicalDevices() local
557 partitioned_output = partitioned_output_user; in RemapOutputsFromLogicalDevices()
561 if (partitioned_output) { in RemapOutputsFromLogicalDevices()
563 return partitioned_output.emitOpError() in RemapOutputsFromLogicalDevices()
574 llvm::enumerate(partitioned_output.output())) { in RemapOutputsFromLogicalDevices()
586 llvm::zip(partitioned_output.output(), tile_sharded_outputs)) in RemapOutputsFromLogicalDevices()
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
Dtpu_resource_partitioning.cc103 auto partitioned_output = builder.create<TF::TPUPartitionedOutputOp>( in PartitionResourceReadsWrites() local
108 llvm::zip(partitioned_input.inputs(), partitioned_output.output())) in PartitionResourceReadsWrites()
Dtpu_sharding_identification_pass.cc193 if (auto partitioned_output = in GetXlaShardingFromResult() local
195 return partitioned_output._XlaSharding(); in GetXlaShardingFromResult()
Dtf_passes.td581 …%partitioned_output:2 = "tf.TPUPartitionedOutput"(%computation) {N = 2 : i64, _XlaSharding = "", p…
582 …"tf.AssignVariableOp"(%arg0, %partitioned_output#0) : (tensor<!tf.resource<tensor<i32>>>, tensor<i…
583 …"tf.AssignVariableOp"(%arg1, %partitioned_output#1) : (tensor<!tf.resource<tensor<i32>>>, tensor<i…
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/tests/
Dtpu_rewrite.mlir1238 …%partitioned_output:2 = "tf.TPUPartitionedOutput"(%computation) {N = 2 : i64, partition_dim = -1 :…
1241 …"tf.AssignVariableOp"(%arg0, %partitioned_output#0) : (tensor<!tf.resource<tensor<i32>>>, tensor<i…
1242 …"tf.AssignVariableOp"(%arg1, %partitioned_output#1) : (tensor<!tf.resource<tensor<i32>>>, tensor<i…
1276 …%partitioned_output:2 = "tf.TPUPartitionedOutput"(%computation) {_XlaSharding = "\08\03\1A\02\01\0…
1279 …"tf.AssignVariableOp"(%arg0, %partitioned_output#0) : (tensor<!tf.resource<tensor<3x2xf32>>>, tens…
1280 …"tf.AssignVariableOp"(%arg1, %partitioned_output#1) : (tensor<!tf.resource<tensor<3x2xf32>>>, tens…
1301 …%partitioned_output:2 = "tf.TPUPartitionedOutput"(%computation) {N = 2 : i64, partition_dim = -1 :…
1302 …"tf.AssignVariableOp"(%arg0, %partitioned_output#0) : (tensor<!tf.resource<tensor<i32>>>, tensor<i…
1303 …"tf.AssignVariableOp"(%arg1, %partitioned_output#1) : (tensor<!tf.resource<tensor<i32>>>, tensor<i…
1324 …%partitioned_output:2 = "tf.TPUPartitionedOutput"(%computation) {N = 2 : i64, partition_dim = -1 :…
[all …]