Lines Matching refs:split_dim
29 static Graph* MakeGraph(int split_dim, int num_split,
33 in_shape.set_dim(split_dim, in_shape.dim_size(split_dim) * num_split);
36 Tensor split_dim_tensor = test::AsScalar<int32>(split_dim);
55 auto g = MakeGraph(/* split_dim = */ 0, num_split, {chunk_size}); \
61 #define BM_SPLIT_2D(split_dim, num_split, chunk_size0, chunk_size1) \
63 BM_Split_2d_##split_dim##_##num_split##_##chunk_size0##_##chunk_size1( \
70 num_split, split_dim, chunk_size0, chunk_size1); \
73 auto g = MakeGraph(split_dim, num_split, {chunk_size0, chunk_size1}); \
78 BM_Split_2d_##split_dim##_##num_split##_##chunk_size0##_##chunk_size1);