- Sort Score
- Result 10 results
- Languages All
Results 1 - 9 of 9 for partition_dim (0.15 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_resource_partitioning.mlir
// CHECK: [[INPUT:%.+]] = "tf.TPUPartitionedInputV2"([[READ0]], [[READ1]]) // CHECK-SAME: _XlaSharding = "" // CHECK-SAME: partition_dims = [] %0 = "tf.TPUPartitionedInputV2"(%arg0, %arg1) {N = 2 : i64, _XlaSharding = "", partition_dims = []} : (tensor<!tf_type.resource<tensor<i32>>>, tensor<!tf_type.resource<tensor<i32>>>) -> tensor<!tf_type.resource<tensor<i32>>>
Registered: Sun Jun 16 05:45:23 UTC 2024 - Last Modified: Tue Jan 23 23:53:20 UTC 2024 - 15.7K bytes - Viewed (0) -
tensorflow/compiler/mlir/tensorflow/tests/tpu_sharding_identification.mlir
%1:2 = "tf_device.cluster_func"(%0, %arg1) {func = @cluster_func, use_spmd_for_xla_partitioning = true, num_cores_per_replica = 1 : i64} : (tensor<*xi32>, tensor<*xi32>) -> (tensor<*xi32>, tensor<*xi32>) %2 = "tf.TPUPartitionedOutputV2"(%1#1) {_XlaSharding = "\08\01\1A\01\01\22\01\00", partition_dims = []} : (tensor<*xi32>) -> tensor<*xi32>
Registered: Sun Jun 16 05:45:23 UTC 2024 - Last Modified: Tue Feb 20 19:07:52 UTC 2024 - 47.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) -
tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir
^bb0(%arg0: tensor<2x4xf32>, %arg1: tensor<2x4xf32>): // expected-error @+1 {{expected 1 inputs, got 2}} %0 = "tf.TPUPartitionedInputV2"(%arg0, %arg1) {partition_dims = [2, 1], is_packed = true} : (tensor<2x4xf32>, tensor<2x4xf32>) -> tensor<4x4xf32> func.return %0 : tensor<4x4xf32> } // ----- // Test invalid tf.TPUPartitionedInputV2 without packing
Registered: Sun Jun 16 05:45:23 UTC 2024 - Last Modified: Mon Oct 23 14:40:35 UTC 2023 - 236.4K bytes - Viewed (0)