Home
last modified time | relevance | path

Searched refs:dim_to_split (Results 1 – 2 of 2) sorted by relevance

/aosp_15_r20/external/pytorch/aten/src/ATen/
H A DTensorIterator.cpp1329 int dim_to_split = -1; in get_dim_to_split() local
1341 dim_to_split = dim; in get_dim_to_split()
1346 return dim_to_split; in get_dim_to_split()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dspmd_partitioner.cc1728 int64_t dim_to_split, int64_t dim_size, in SplitReshapeHelper() argument
1733 shape_dim.insert(shape_dim.begin() + dim_to_split + 1, 1); in SplitReshapeHelper()
1737 base_shape_dim.insert(base_shape_dim.begin() + dim_to_split + 1, dim_size); in SplitReshapeHelper()
1738 base_shape_dim[dim_to_split] /= dim_size; in SplitReshapeHelper()