| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/ops/ |
| H A D | tpu_embedding_shape_util.cc | 38 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()
|
| H A D | tpu_embedding_ops.cc | 238 int num_shards; in __anoncc2b8c080702() local
321 int num_shards; in __anoncc2b8c080802() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/data/kernel_tests/ |
| H A D | shard_test.py | 104 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
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/kernel_tests/nn_ops/ |
| H A D | embedding_ops_test.py | 135 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
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| H A D | combine_tpu_embedding_load_retrieve_pass.cc | 189 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
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/kernels/data/ |
| H A D | shard_dataset_op.cc | 45 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
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/ |
| H A D | tpu_embedding_load_retrieve_ops.cc | 49 int shard_id, int num_shards, in ComputeExpectedTableShardShapes()
110 int num_shards; in LoadAllTPUEmbeddingParametersOp() local
253 int num_shards; in RetrieveAllTPUEmbeddingParametersOp() local
|
| /aosp_15_r20/external/rust/beto-rust/nearby/presence/np_ffi_core/src/ |
| D | common.rs | 42 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
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/ |
| H A D | tpu_embedding_optimization_parameters_utils.cc | 425 int num_shards; in operator ()() local
455 int num_shards; in operator ()() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/distribute/v1/ |
| H A D | all_reduce_test.py | 185 def _buildShuffle(self, num_workers, num_gpus, num_shards): argument
192 def _testShuffleAllReduce(self, num_workers, num_gpus, shape, num_shards): argument
|
| /aosp_15_r20/external/pytorch/torch/csrc/distributed/autograd/context/ |
| H A D | container.cpp | 23 DistAutogradContainer::DistAutogradContainer(uint32_t num_shards) in DistAutogradContainer()
69 uint32_t num_shards = 1; in computeNumShards() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/kernels/ |
| H A D | range_sampler.cc | 240 int32_t num_shards, int32_t shard) in FixedUnigramSampler()
256 int32_t num_shards, int32_t shard) in FixedUnigramSampler()
|
| H A D | topk_op_gpu.h | 263 __device__ void mergeShards(int num_shards, int k, in mergeShards()
374 cudaError LaunchTopKKernel(const gpuStream_t& stream, int num_shards, in LaunchTopKKernel()
|
| /aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/ |
| H A D | meta_spmd_expander.cc | 648 const int num_shards = input_layout.num_shards_for_dim( in MakeLayoutForReshape() local
1007 const int num_shards = output_layout->num_shards()[axis]; in ExpandOp() local
1196 auto num_shards = in ExpandOp() local
|
| H A D | fill_spmd_expander.cc | 73 auto num_shards = in ExpandOp() local
|
| H A D | nullary_spmd_expander.cc | 66 const int num_shards = in ExpandOp() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/lib/core/ |
| H A D | threadpool_test.cc | 69 int64_t num_shards = 0; in RunWithFixedBlockSize() local
121 int64_t num_shards = 0; in RunWithFixedBlockSizeTransformRangeConcurrently() local
|
| H A D | blocking_counter_test.cc | 57 const int num_shards = num_threads * shards_per_thread; in BM_BlockingCounter() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/util/tensor_bundle/ |
| H A D | naming.cc | 28 string DataFilename(StringPiece prefix, int32_t shard_id, int32_t num_shards) { in DataFilename()
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/distribute/ |
| H A D | input_ops.py | 26 def auto_shard_dataset(dataset, num_shards, index, num_replicas_in_sync=None): argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/protobuf/ |
| H A D | tensor_bundle.proto | 27 int32 num_shards = 1; field
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/util/ |
| H A D | work_sharder.cc | 67 const int num_shards = in Do() local
|
| H A D | work_sharder_test.cc | 34 int64_t num_shards = 0; in RunSharding() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/tpu/ |
| H A D | tpu_optimizer.py | 73 def _verify_and_get_subgroup_size(self, group_assignment, num_shards): argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/ops/ |
| H A D | partitioned_variables.py | 221 def fixed_size_partitioner(num_shards, axis=0): argument
|