- Sort Score
- Result 10 results
- Languages All
Results 1 - 6 of 6 for partition_dim (0.18 sec)
-
tensorflow/compiler/mlir/tensorflow/transforms/tpu_partitioned_op_conversion.cc
int rank = tensor_type.getRank(); if (rank <= partition_dim) { return op->emitError() << "cannot partition " << first_operand_type << " (rank = " << rank << ") along dimension " << partition_dim << "."; } llvm::SmallVector<int64_t, 4> partition_dims(is_replicated ? 0 : rank, 1); if (!is_replicated) { partition_dims[partition_dim] = num_cores_per_replica.getInt(); }
Registered: Sun Jun 16 05:45:23 UTC 2024 - Last Modified: Thu Apr 25 16:01:03 UTC 2024 - 5.1K bytes - Viewed (0) -
tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.td
%pi_0 = "tf.TPUPartitionedInput"(%arg0, %arg1) {_XlaSharding = "", device = "", partition_dim = -1 : i64} : (!rtype, !rtype) -> !rtype %pi_1 = "tf.TPUPartitionedInput"(%arg2, %arg3) {_XlaSharding = "", device = "", partition_dim = -1 : i64} : (!rtype, !rtype) -> !rtype %ri = "tf.TPUReplicatedInput"(%pi_0, %pi_1) : (!rtype, !rtype) -> !rtype return %ri : !rtype
Registered: Sun Jun 16 05:45:23 UTC 2024 - Last Modified: Wed Jun 12 21:18:05 UTC 2024 - 99.6K bytes - Viewed (0) -
tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td
let summary = [{ An op that groups a list of partitioned inputs together. This op }]; let arguments = (ins Variadic<TF_Tensor>:$inputs, DefaultValuedOptionalAttr<I64Attr, "0">:$partition_dim, OptionalAttr<StrAttr>:$_XlaSharding ); let results = (outs TF_Tensor:$output ); TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
Registered: Sun Jun 16 05:45:23 UTC 2024 - Last Modified: Wed Apr 24 04:08:35 UTC 2024 - 90.5K bytes - Viewed (0) -
tensorflow/compiler/mlir/tensorflow/tests/tpu_cluster_formation.mlir
func.func @partitioned_outputs(%arg0: tensor<?xi32>) -> (tensor<?xi32>, tensor<?xi32>, tensor<?xi32>, tensor<?xi32>) { %pi0 = "tf.TPUPartitionedInputV2"(%arg0) {N = 2, partition_dims = [], _XlaSharding = "", is_packed = true} : (tensor<?xi32>) -> (tensor<?xi32>) %pi1 = "tf.TPUPartitionedInputV2"(%arg0) {N = 2, partition_dims = [], _XlaSharding = "", is_packed = true} : (tensor<?xi32>) -> (tensor<?xi32>)
Registered: Sun Jun 16 05:45:23 UTC 2024 - Last Modified: Thu May 02 22:03:30 UTC 2024 - 53.9K bytes - Viewed (0) -
tensorflow/compiler/mlir/tensorflow/tests/tpu_rewrite.mlir
%read1 = "tf.ReadVariableOp"(%arg1) : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32> // CHECK-NOT: tf.TPUPartitionedInputV2 %partitioned_input = "tf.TPUPartitionedInputV2"(%read0, %read1) {N = 2 : i64, partition_dims = []} : (tensor<i32>, tensor<i32>) -> tensor<i32> // CHECK: %[[COMPILE_OUTPUT:[0-9]*]]:3 = "tf_device.launch" // CHECK-NEXT: "tf._TPUCompileMlir"() // CHECK: "tf_device.launch"
Registered: Sun Jun 16 05:45:23 UTC 2024 - Last Modified: Thu May 02 22:03:30 UTC 2024 - 172.9K bytes - Viewed (0) -
tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc
TPUPartitionedInputV2Op op = *this; int num_partitions = 1; const mlir::ArrayAttr partition_dims = op.getPartitionDims(); for (const mlir::Attribute &dim : partition_dims) { num_partitions *= dim.cast<IntegerAttr>().getInt(); } const bool is_packed = op.getIsPacked(); const bool replicated = partition_dims.empty(); const int num_inputs_expected = is_packed ? 1 : num_partitions;
Registered: Sun Jun 16 05:45:23 UTC 2024 - Last Modified: Thu May 09 22:07:10 UTC 2024 - 170.8K bytes - Viewed (0)