Home
last modified time | relevance | path

Searched refs:split_factor (Results 1 – 8 of 8) sorted by relevance

/aosp_15_r20/external/pytorch/test/distributed/_tensor/
H A Dtest_utils.py135 placements = [_StridedShard(0, split_factor=tp_size), Shard(0)]
162 placements = [Replicate(), _StridedShard(0, split_factor=tp_size), Shard(0)]
197 _StridedShard(shard_dim, split_factor=mesh_size_1),
217 _StridedShard(0, split_factor=mesh_size_1),
228 _StridedShard(0, split_factor=mesh_size_1),
250 _StridedShard(0, split_factor=mesh_size_1),
251 _StridedShard(1, split_factor=mesh_size_3),
267 _StridedShard(0, split_factor=mesh_size_1),
268 _StridedShard(1, split_factor=mesh_size_3),
291 shard_placement = _StridedShard(0, split_factor=1) # same as Shard(0)
[all …]
/aosp_15_r20/external/pytorch/torch/distributed/tensor/
H A Dplacement_types.py419 split_factor: int
423 return self.dim == other.dim and self.split_factor == other.split_factor
432 return hash((self.dim, self.split_factor))
459 total_split = num_chunks * self.split_factor
466 group_size = self.split_factor
495 total_split = num_chunks * self.split_factor
521 idx_after_split = idx % num_chunks * self.split_factor + idx // num_chunks
H A D_utils.py178 placement.split_factor * mesh_dim_size
/aosp_15_r20/external/pytorch/torch/distributed/_composable/fsdp/
H A D_fsdp_param.py279 split_factor = self._tp_spec.num_shards_map[0]
286 _StridedShard(0, split_factor=split_factor)
287 if split_factor > 1
304 if split_factor > 1: # FSDP has strided sharding on tensor dim 0
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/
H A Dcudnn_vectorize_convolutions.cc149 int64_t split_factor = vect_size / size(vect_dim); in RevectorizeInstr() local
150 CHECK_EQ(size(dim) % split_factor, 0); in RevectorizeInstr()
153 instr = SplitAtDim(instr, dim, split_factor); in RevectorizeInstr()
/aosp_15_r20/external/mesa3d/src/intel/vulkan/
H A DgenX_acceleration_structure.c212 const float split_factor = 0.0f; in get_gpu_size_estimate() local
215 num_prims_to_split = num_triangles + (double)split_factor; in get_gpu_size_estimate()
/aosp_15_r20/external/pytorch/test/distributed/_composable/fsdp/
H A Dtest_fully_shard_init.py422 _StridedShard(0, split_factor=tp_mesh.size()),
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dspmd_partitioner.cc2707 int64_t split_factor = input_dim_size / output_dim_size; in HandleReshape() local
2721 dim->set_padding_high(output_shard_size * split_factor * in HandleReshape()
2739 output_shard_size * split_factor); in HandleReshape()