| /external/tensorflow/tensorflow/core/kernels/data/ |
| D | shard_dataset_op.cc | 44 Dataset(OpKernelContext* ctx, int64_t num_shards, int64_t index, in Dataset() argument
47 num_shards_(num_shards), in Dataset()
54 strings::Printf("%lld", static_cast<long long>(num_shards))}}) { in Dataset()
103 Node* num_shards = nullptr; in AsGraphDefInternal() local
104 TF_RETURN_IF_ERROR(b->AddScalar(num_shards_, &num_shards)); in AsGraphDefInternal()
112 b->AddDataset(this, {input_graph_node, num_shards, index}, in AsGraphDefInternal()
257 int64_t num_shards = 0; in MakeDataset() local
259 OP_REQUIRES_OK(ctx, ParseScalarArgument<int64>(ctx, kNumShards, &num_shards)); in MakeDataset()
261 ctx, num_shards > 0 || num_shards == kShardHint, in MakeDataset()
264 num_shards, ").")); in MakeDataset()
[all …]
|
| /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 | 77 def _verify_and_get_subgroup_size(self, group_assignment, num_shards): argument
104 if set(range(num_shards)) != replica_ids:
149 num_shards = tpu_function.get_tpu_context().number_of_shards
150 if num_shards is None:
154 num_shards = 1
157 num_shards)
159 if num_shards > 1 and self._reduction == losses.Reduction.MEAN:
163 scale = 1.0 / num_shards
|
| /external/tensorflow/tensorflow/core/kernels/ |
| D | topk_op_gpu.h | 262 __device__ void mergeShards(int num_shards, int k, in mergeShards() argument
270 const int heap_size = k < num_shards ? k : num_shards; in mergeShards()
284 for (int shard = heap_size; shard < num_shards; shard++) { in mergeShards()
316 int next_shard_index = shard_index + num_shards; in mergeShards()
373 cudaError LaunchTopKKernel(const gpuStream_t& stream, int num_shards, in LaunchTopKKernel() argument
383 if (num_shards <= 0) { in LaunchTopKKernel()
387 num_shards = shared_memory_size / heap_size - 1; in LaunchTopKKernel()
388 if (num_shards <= 0) { in LaunchTopKKernel()
389 num_shards = 1; in LaunchTopKKernel()
391 auto shard_size = length / num_shards; in LaunchTopKKernel()
[all …]
|
| /external/tensorflow/tensorflow/core/util/tensor_bundle/ |
| D | naming.cc | 28 string DataFilename(StringPiece prefix, int32_t shard_id, int32_t 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_t num_shards = 0; in RunSharding() local
38 [=, &mu, &num_shards, &num_done_work, &work](int64_t start, in RunSharding()
44 ++num_shards; in RunSharding()
52 << num_shards; in RunSharding()
61 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_t block_size = (total + num_shards - 1) / num_shards; in Do()
|
| /external/tensorflow/tensorflow/python/data/kernel_tests/ |
| D | shard_test.py | 102 def _build_dataset(self, num_elements, num_shards, index): argument
103 return dataset_ops.Dataset.range(num_elements).shard(num_shards, index)
110 elems=[10, 100], num_shards=[2, 5], index=[0, 1])))
111 def test(self, verify_fn, elems, num_shards, index): argument
114 lambda: self._build_dataset(elems, num_shards, index),
115 num_outputs=elems // num_shards)
|
| /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 | LoadTPUEmbeddingMomentumParameters.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 | RetrieveTPUEmbeddingProximalAdagradParameters.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 | RetrieveTPUEmbeddingMomentumParameters.pbtxt | 28 name: "num_shards"
64 name: "num_shards"
105 name: "num_shards"
|
| D | LoadTPUEmbeddingAdadeltaParameters.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 | LoadTPUEmbeddingMDLAdagradLightParameters.pbtxt | 36 name: "num_shards"
80 name: "num_shards"
129 name: "num_shards"
|
| D | RetrieveTPUEmbeddingADAMParameters.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 | LoadTPUEmbeddingADAMParameters.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 | RetrieveTPUEmbeddingAdadeltaParameters.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"
|