Searched refs:split_factor (Results 1 – 8 of 8) sorted by relevance
/aosp_15_r20/external/pytorch/test/distributed/_tensor/ |
H A D | test_utils.py | 135 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 D | placement_types.py | 419 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.py | 178 placement.split_factor * mesh_dim_size
|
/aosp_15_r20/external/pytorch/torch/distributed/_composable/fsdp/ |
H A D | _fsdp_param.py | 279 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 D | cudnn_vectorize_convolutions.cc | 149 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 D | genX_acceleration_structure.c | 212 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 D | test_fully_shard_init.py | 422 _StridedShard(0, split_factor=tp_mesh.size()),
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
H A D | spmd_partitioner.cc | 2707 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()
|