Home
last modified time | relevance | path

Searched defs:num_shards (Results 1 – 25 of 63) sorted by relevance

123

/external/tensorflow/tensorflow/core/tpu/ops/
Dtpu_embedding_shape_util.cc38 int64 vocabulary_size, int table_dimension, int shard_id, int num_shards, in ComputeOneTableShape()
64 const absl::Span<const int> table_dimensions, int shard_id, int num_shards, in ComputeTableShapes()
76 const TPUEmbeddingConfiguration& config, int shard_id, int num_shards, in ComputeTableShapes()
98 int num_shards) { in ComputeNumEmbeddingIdsPerShard()
Dtpu_embedding_ops.cc238 int num_shards; in __anoned4aa6880702() local
321 int num_shards; in __anoned4aa6880802() local
/external/tensorflow/tensorflow/python/data/kernel_tests/
Dshard_test.py104 def _build_dataset(self, num_elements, num_shards, index): argument
113 def test(self, verify_fn, elems, num_shards, index): argument
166 def testMultipleCombinations(self, elements, num_shards, index): argument
/external/tensorflow/tensorflow/core/kernels/data/
Dshard_dataset_op.cc45 Dataset(OpKernelContext* ctx, int64_t num_shards, int64_t index, in Dataset()
118 Node* num_shards = nullptr; in AsGraphDefInternal() local
272 int64_t num_shards = 0; in MakeDataset() local
/external/tensorflow/tensorflow/core/util/tensor_bundle/
Dnaming.cc28 string DataFilename(StringPiece prefix, int32_t shard_id, int32_t num_shards) { in DataFilename()
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_embedding_load_retrieve_ops.cc49 int shard_id, int num_shards, in ComputeExpectedTableShardShapes()
110 int num_shards; in LoadAllTPUEmbeddingParametersOp() local
253 int num_shards; in RetrieveAllTPUEmbeddingParametersOp() local
/external/tensorflow/tensorflow/core/tpu/
Dtpu_embedding_optimization_parameters_utils.cc425 int num_shards; in operator ()() local
455 int num_shards; in operator ()() local
/external/tensorflow/tensorflow/python/kernel_tests/nn_ops/
Dembedding_ops_test.py135 def _EmbeddingParams(num_shards, argument
165 def _EmbeddingParamsAsPartitionedVariable(num_shards, argument
185 num_shards, argument
812 def _random_weights(self, vocab_size=4, embed_dim=4, num_shards=1): argument
/external/tensorflow/tensorflow/python/distribute/
Dinput_ops.py26 def auto_shard_dataset(dataset, num_shards, index, num_replicas_in_sync=None): argument
/external/rust/beto-rust/nearby/presence/np_ffi_core/src/
Dcommon.rs42 num_shards: u8, field
56 pub(crate) fn num_shards(&self) -> u8 { in num_shards() method
67 pub(crate) fn num_shards(&self) -> u8 { in num_shards() method
/external/tensorflow/tensorflow/python/distribute/v1/
Dall_reduce_test.py185 def _buildShuffle(self, num_workers, num_gpus, num_shards): argument
192 def _testShuffleAllReduce(self, num_workers, num_gpus, shape, num_shards): argument
/external/tensorflow/tensorflow/core/protobuf/
Dtensor_bundle.proto27 int32 num_shards = 1; field
/external/tensorflow/tensorflow/core/tpu/graph_rewrite/
Dcombine_tpu_embedding_load_retrieve_pass.cc189 std::vector<bool>* is_debug_load_retrieve_node, int* num_shards) { in GetLoadOrRetrieveNodesByTable()
282 LoadCombinedParametersType* combined_inputs, int* num_shards) { in CombinePerTableParametersForLoad()
475 int num_shards; in Run() local
645 int num_shards; in Run() local
/external/tensorflow/tensorflow/dtensor/mlir/expansions/
Dfill_spmd_expander.cc73 auto num_shards = in ExpandOp() local
Dnullary_spmd_expander.cc66 const int num_shards = in ExpandOp() local
/external/pytorch/torch/csrc/distributed/autograd/context/
Dcontainer.cpp23 DistAutogradContainer::DistAutogradContainer(uint32_t num_shards) in DistAutogradContainer()
69 uint32_t num_shards = 1; in computeNumShards() local
/external/tensorflow/tensorflow/core/util/
Dwork_sharder.cc67 const int num_shards = in Do() local
Dwork_sharder_test.cc34 int64_t num_shards = 0; in RunSharding() local
/external/tensorflow/tensorflow/core/lib/core/
Dblocking_counter_test.cc57 const int num_shards = num_threads * shards_per_thread; in BM_BlockingCounter() local
/external/tensorflow/tensorflow/python/tpu/
Dtpu_optimizer.py73 def _verify_and_get_subgroup_size(self, group_assignment, num_shards): argument
/external/tensorflow/tensorflow/python/ops/
Dpartitioned_variables.py221 def fixed_size_partitioner(num_shards, axis=0): argument
Dcandidate_sampling_ops.py228 num_shards=1, argument
/external/pytorch/torch/distributed/tensor/
D_dtensor_spec.py122 def num_shards(self) -> int: member in DTensorSpec
/external/tensorflow/tensorflow/core/kernels/
Drange_sampler.cc240 int32_t num_shards, int32_t shard) in FixedUnigramSampler()
256 int32_t num_shards, int32_t shard) in FixedUnigramSampler()
/external/tensorflow/tensorflow/dtensor/mlir/ir/
Dtf_dtensor.cc60 const auto& num_shards = layout.num_shards(); in verify() local

123