| /external/tensorflow/tensorflow/core/kernels/ |
| D | dynamic_partition_op.cc | 34 OP_REQUIRES_OK(c, c->GetAttr("num_partitions", &num_partitions_)); in DynamicPartitionOp_Shared()
56 gtl::InlinedVector<int, 32> partition_count(num_partitions_); in ValidateAndAllocateOutputs()
61 OP_REQUIRES(c, FastBoundsCheck(p, num_partitions_), in ValidateAndAllocateOutputs()
64 " = ", p, " is not in [0, ", num_partitions_, ")")); in ValidateAndAllocateOutputs()
70 for (int p = 0; p < num_partitions_; p++) { in ValidateAndAllocateOutputs()
82 int num_partitions_; member in tensorflow::DynamicPartitionOp_Shared
96 if (num_partitions_ == 0 || data->NumElements() == 0) return; in Compute()
100 gtl::InlinedVector<int, 32> output_index(num_partitions_); in Compute()
108 out_vec.reserve(num_partitions_); in Compute()
109 for (int p = 0; p < num_partitions_; p++) { in Compute()
[all …]
|
| D | dynamic_partition_op_gpu.cu.cc | 199 OP_REQUIRES_OK(c, c->GetAttr("num_partitions", &num_partitions_)); in DynamicPartitionOpGPU()
200 OP_REQUIRES(c, num_partitions_ >= 1, in DynamicPartitionOpGPU()
207 int32 M = std::max(N, num_partitions_); in AllocateTempSpace()
224 for (int p = 0; p < num_partitions_; p++) { in AllocateOutputs()
256 c->allocate_temp(DT_INT32, TensorShape({num_partitions_}), in ComputeAsync()
260 for (int i = 0; i < num_partitions_; i++) e_part_count(i) = 0; in ComputeAsync()
271 c->allocate_temp(DT_INT32, TensorShape({num_partitions_}), in ComputeAsync()
297 num_partitions_ * sizeof(int32)); in ComputeAsync()
301 num_partitions_ * sizeof(int32)) in ComputeAsync()
394 c->allocate_temp(DT_INT32, TensorShape({num_partitions_}), in CountAndSortParts()
[all …]
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | fft_handler.cc | 359 if (input_length != fft_length || input_length % num_partitions_ != 0) { in HandleFft()
366 num_partitions_) { in HandleFft()
384 partitioned_input.hlo(), num_partitions_, hlo->sharding(), in HandleFft()
396 result = ShuffleWithinEachPartitionUsingOneHot(result, num_partitions_, in HandleFft()
403 result, num_partitions_, partitioned_input.state().collective_ops_creator, in HandleFft()
410 partitioned_fft_length.back() /= num_partitions_; in HandleFft()
416 result, num_partitions_, partitioned_input.state().partition_id, in HandleFft()
424 num_partitions_, partitioned_input.state().partition_id, in HandleFft()
|
| D | spmd_partitioner.h | 189 : num_partitions_(num_partitions), in SpmdPartitioner()
246 const int64 num_partitions_; variable
524 int64 num_partitions_; variable
|
| D | spmd_partitioner.cc | 1358 num_partitions_(num_partitions), in SpmdPartitioningVisitor()
2064 num_partitions_ / replication_count - in HandleReshape()
2103 num_partitions_ / replication_count); in HandleReshape()
2120 num_partitions_ / replication_count); in HandleReshape()
2522 std::vector<int32> conditional_branch_indices(num_partitions_); in HandleInfeed()
2523 for (int64 i = 0; i < num_partitions_; ++i) { in HandleInfeed()
2544 if (per_branch_partitioned_shapes.size() == num_partitions_) { in HandleInfeed()
2924 std::vector<int32> conditional_branch_indices(num_partitions_); in HandleOutfeed()
2925 for (int64 i = 0; i < num_partitions_; ++i) { in HandleOutfeed()
2947 if (per_branch_partitioned_shapes.size() == num_partitions_) { in HandleOutfeed()
[all …]
|
| D | dot_handler.cc | 2850 num_partitions_, create_sharded_dot, conv_window, module_, in HandleDotHelper()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_module_config.h | 132 num_partitions_ = num_partitions; in set_num_partitions()
134 int64 num_partitions() const { return num_partitions_; } in num_partitions()
257 int64 num_partitions_ = 1; variable
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | executable_build_options.h | 84 int num_partitions() const { return num_partitions_; } in num_partitions()
148 int num_partitions_ = 1; variable
|
| D | executable_build_options.cc | 70 num_partitions_ = num_partitions; in set_num_partitions()
|
| /external/tensorflow/tensorflow/lite/delegates/ |
| D | utils.cc | 56 int num_partitions_ = 0; in Partition() local
59 &num_partitions_) != kTfLiteOk) { in Partition()
64 for (int i = 0; i < num_partitions_; ++i) { in Partition()
|