| /external/tensorflow/tensorflow/python/kernel_tests/ |
| D | embedding_ops_test.py | 140 def _EmbeddingParams(num_shards, argument
150 for i in range(num_shards):
151 shard_shape = [vocab_size // num_shards] + shape
152 if i < vocab_size % num_shards: # Excess goes evenly on the first shards
170 def _EmbeddingParamsAsPartitionedVariable(num_shards, argument
176 num_shards, vocab_size, dtype=dtype, shape=shape)
183 max_partitions=num_shards, min_slice_size=1),
190 num_shards, argument
209 val = np.copy(params[_PName(i % num_shards) + ":0"][
210 i // num_shards, :]) * weight_value
[all …]
|
| /external/tensorflow/tensorflow/python/tpu/ |
| D | tpu_optimizer.py | 72 def _verify_and_get_subgroup_size(self, group_assignment, num_shards): argument
98 if set(range(num_shards)) != replica_ids:
101 num_shards, group_assignment))
144 num_shards = tpu_function.get_tpu_context().number_of_shards
145 if num_shards is None:
149 num_shards = 1
152 num_shards)
154 if num_shards > 1 and self._reduction == losses.Reduction.MEAN:
158 scale = 1.0 / num_shards
|
| /external/tensorflow/tensorflow/python/data/experimental/kernel_tests/serialization/ |
| D | shard_dataset_serialization_test.py | 33 def _build_dataset(self, num_elements, num_shards, index): argument
34 return dataset_ops.Dataset.range(num_elements).shard(num_shards, index)
40 elems=[10, 100], num_shards=[2, 5], index=[0, 1])))
41 def testCore(self, elems, num_shards, index): argument
42 self.run_core_tests(lambda: self._build_dataset(elems, num_shards, index),
43 elems // num_shards)
|
| /external/tensorflow/tensorflow/core/kernels/data/ |
| D | shard_dataset_op.cc | 43 Dataset(OpKernelContext* ctx, int64 num_shards, int64 index, in Dataset() argument
46 num_shards_(num_shards), in Dataset()
53 strings::Printf("%lld", static_cast<long long>(num_shards))}}) { in Dataset()
102 Node* num_shards = nullptr; in AsGraphDefInternal() local
103 TF_RETURN_IF_ERROR(b->AddScalar(num_shards_, &num_shards)); in AsGraphDefInternal()
111 b->AddDataset(this, {input_graph_node, num_shards, index}, in AsGraphDefInternal()
247 int64 num_shards = 0; in MakeDataset() local
249 OP_REQUIRES_OK(ctx, ParseScalarArgument<int64>(ctx, kNumShards, &num_shards)); in MakeDataset()
251 ctx, num_shards > 0, in MakeDataset()
254 num_shards, ").")); in MakeDataset()
[all …]
|
| /external/tensorflow/tensorflow/core/kernels/ |
| D | topk_op_gpu.h | 271 __device__ void mergeShards(int num_shards, int k,
279 const int heap_size = k < num_shards ? k : num_shards;
293 for (int shard = heap_size; shard < num_shards; shard++) {
325 int next_shard_index = shard_index + num_shards;
382 cudaError LaunchTopKKernel(const gpuStream_t& stream, int num_shards,
392 if (num_shards <= 0) {
396 num_shards = shared_memory_size / heap_size - 1;
397 if (num_shards <= 0) {
398 num_shards = 1;
400 auto shard_size = length / num_shards;
[all …]
|
| /external/tensorflow/tensorflow/core/util/tensor_bundle/ |
| D | naming.cc | 28 string DataFilename(StringPiece prefix, int32 shard_id, int32 num_shards) { in DataFilename() argument
29 DCHECK_GT(num_shards, 0); in DataFilename()
30 DCHECK_LT(shard_id, num_shards); in DataFilename()
33 shard_id, num_shards); in DataFilename()
|
| /external/tensorflow/tensorflow/core/util/ |
| D | work_sharder_test.cc | 34 int64 num_shards = 0; in RunSharding() local
38 [=, &mu, &num_shards, &num_done_work, &work](int64 start, int64 limit) { in RunSharding()
43 ++num_shards; in RunSharding()
51 << num_shards; in RunSharding()
60 EXPECT_LE(num_shards, 1 + per_thread_max_parallelism); in RunSharding()
|
| D | work_sharder.cc | 67 const int num_shards = in Do() local
77 const int64 block_size = (total + num_shards - 1) / num_shards; in Do()
|
| /external/tensorflow/tensorflow/core/ops/compat/ops_history_v1/ |
| D | LoadTPUEmbeddingStochasticGradientDescentParameters.pbtxt | 24 name: "num_shards"
56 name: "num_shards"
93 name: "num_shards"
|
| D | RetrieveTPUEmbeddingStochasticGradientDescentParameters.pbtxt | 24 name: "num_shards"
56 name: "num_shards"
93 name: "num_shards"
|
| D | LoadTPUEmbeddingAdagradParameters.pbtxt | 28 name: "num_shards"
64 name: "num_shards"
105 name: "num_shards"
|
| D | RetrieveTPUEmbeddingProximalAdagradParameters.pbtxt | 28 name: "num_shards"
64 name: "num_shards"
105 name: "num_shards"
|
| D | RetrieveTPUEmbeddingMomentumParameters.pbtxt | 28 name: "num_shards"
64 name: "num_shards"
105 name: "num_shards"
|
| D | RetrieveTPUEmbeddingAdagradParameters.pbtxt | 28 name: "num_shards"
64 name: "num_shards"
105 name: "num_shards"
|
| D | LoadTPUEmbeddingMomentumParameters.pbtxt | 28 name: "num_shards"
64 name: "num_shards"
105 name: "num_shards"
|
| D | LoadTPUEmbeddingProximalAdagradParameters.pbtxt | 28 name: "num_shards"
64 name: "num_shards"
105 name: "num_shards"
|
| D | LoadTPUEmbeddingRMSPropParameters.pbtxt | 32 name: "num_shards"
72 name: "num_shards"
117 name: "num_shards"
|
| D | RetrieveTPUEmbeddingRMSPropParameters.pbtxt | 32 name: "num_shards"
72 name: "num_shards"
117 name: "num_shards"
|
| D | LoadTPUEmbeddingAdadeltaParameters.pbtxt | 32 name: "num_shards"
72 name: "num_shards"
117 name: "num_shards"
|
| D | LoadTPUEmbeddingFTRLParameters.pbtxt | 32 name: "num_shards"
72 name: "num_shards"
117 name: "num_shards"
|
| D | LoadTPUEmbeddingMDLAdagradLightParameters.pbtxt | 36 name: "num_shards"
80 name: "num_shards"
129 name: "num_shards"
|
| D | LoadTPUEmbeddingAdagradParametersGradAccumDebug.pbtxt | 32 name: "num_shards"
72 name: "num_shards"
117 name: "num_shards"
|
| D | RetrieveTPUEmbeddingFTRLParameters.pbtxt | 32 name: "num_shards"
72 name: "num_shards"
117 name: "num_shards"
|
| D | LoadTPUEmbeddingCenteredRMSPropParameters.pbtxt | 36 name: "num_shards"
80 name: "num_shards"
129 name: "num_shards"
|
| D | RetrieveTPUEmbeddingProximalAdagradParametersGradAccumDebug.pbtxt | 32 name: "num_shards"
72 name: "num_shards"
117 name: "num_shards"
|