Home
last modified time | relevance | path

Searched refs:num_split (Results 1 – 25 of 83) sorted by relevance

1234

/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/
H A Dsparse_split_op_gpu.cu.cc86 SliceIndexer(const Index split_dim_size, const Index num_split) in SliceIndexer()
87 : split_size_(split_dim_size / num_split), in SliceIndexer()
88 residual_(split_dim_size % num_split) {} in SliceIndexer()
119 Index input_nnz, int num_split, in LaunchSparseSplitSliceIndexesKernel() argument
136 Index input_nnz, int num_split, in SparseSplitFindSliceEndsKernel() argument
139 for (int slice_index : GpuGridRangeX<int>(num_split)) { in SparseSplitFindSliceEndsKernel()
147 Index input_nnz, int num_split, in LaunchSparseSplitFindSliceEndsKernel() argument
151 num_split, device, &SparseSplitFindSliceEndsKernel<Index>, in LaunchSparseSplitFindSliceEndsKernel()
155 device.stream(), input_nnz, num_split, in LaunchSparseSplitFindSliceEndsKernel()
214 const int64_t axis, const int num_split, in operator ()()
[all …]
H A Dsplit_op_test.cc29 static Graph* MakeGraph(int split_dim, int num_split, in MakeGraph() argument
33 in_shape.set_dim(split_dim, in_shape.dim_size(split_dim) * num_split); in MakeGraph()
41 .Attr("num_split", num_split) in MakeGraph()
46 #define BM_SPLIT_1D(num_split, chunk_size) \ argument
47 static void BM_Split_1d_##num_split##_##chunk_size( \
50 strings::Printf("1-D %d chunks of %d each", num_split, chunk_size); \
52 auto g = MakeGraph(/* split_dim = */ 0, num_split, {chunk_size}); \
55 num_split * chunk_size); \
57 BENCHMARK(BM_Split_1d_##num_split##_##chunk_size)->UseRealTime();
59 #define BM_SPLIT_2D(split_dim, num_split, chunk_size0, chunk_size1) \ argument
[all …]
H A Dsplit_op.cc58 const int32_t num_split = num_outputs(); in ComputeEasyCases() local
67 context, num_split > 0, in ComputeEasyCases()
69 "Number of ways to split should be > 0, but got ", num_split)); in ComputeEasyCases()
71 OP_REQUIRES(context, input_shape.dim_size(split_dim) % num_split == 0, in ComputeEasyCases()
76 ") ", "and num_split ", num_split)); in ComputeEasyCases()
78 if (num_split == 1) { in ComputeEasyCases()
95 const int64_t delta = input_shape.dim_size(0) / num_split; in ComputeEasyCases()
96 for (int i = 0; i < num_split; ++i) { in ComputeEasyCases()
138 const ReshapeResultType& reshape_result, int32_t num_split, in operator ()() argument
145 (num_split >= 4 && in operator ()()
[all …]
H A Dsplit_v_op_test.cc78 #define BM_SPLITV_1D(num_split, total_size) \ argument
79 static void BM_SplitV_1d_##num_split##_##total_size( \
82 strings::Printf("1-D %d chunks totaling %d", num_split, total_size); \
85 GenerateRandomIntsWithSum(total_size, num_split), \
91 BENCHMARK(BM_SplitV_1d_##num_split##_##total_size)->UseRealTime();
93 #define BM_SPLITV_2D(split_dim, num_split, dim0, dim1) \ argument
94 static void BM_SplitV_2d_##split_dim##_##num_split##dim0##dim1( \
98 num_split, split_dim, dim0, dim1); \
102 GenerateRandomIntsWithSum(total_size_vec[split_dim], num_split), \
108 BENCHMARK(BM_SplitV_2d_##split_dim##_##num_split##dim0##dim1)->UseRealTime();
[all …]
H A Dsplit_v_op.cc55 const int32_t num_split = context->num_outputs(); in ComputeEasyCases() local
71 split_tensor.dims() == 1 && split_tensor.NumElements() == num_split, in ComputeEasyCases()
85 context, num_split > 0, in ComputeEasyCases()
87 "Number of ways to split should be > 0, but got ", num_split)); in ComputeEasyCases()
98 if (num_split == 1) { in ComputeEasyCases()
161 for (int i = 0; i < num_split; ++i) { in ComputeEasyCases()
238 const int num_split = split_start_points.size(); in operator ()() local
240 (num_split >= kMinimumSplitNum && in operator ()()
241 input_element_count >= std::min(num_threads, num_split) * 4096 && in operator ()()
242 input_element_count < num_split * 180 * 1024); in operator ()()
[all …]
H A Dsparse_split_op.cc36 const int64_t axis, const int num_split, in operator ()()
46 sparse_tensor, axis, num_split, &outputs)); in operator ()()
48 for (int slice_index = 0; slice_index < num_split; ++slice_index) { in operator ()()
50 context->set_output(slice_index + num_split, in operator ()()
54 slice_index + 2 * num_split, in operator ()()
69 void SparseSplitOpImpl(OpKernelContext* context, int num_split, in SparseSplitOpImpl() argument
114 context, num_split >= 1 && num_split <= input_shape.vec<int64_t>()(axis), in SparseSplitOpImpl()
118 num_split), in SparseSplitOpImpl()
130 dense_shape, axis, num_split, done); in SparseSplitOpImpl()
H A Dsplit_lib_gpu.cu.cc77 const int32 num_split = output_ptr_data.size; in SplitOpKernel() local
82 eigen_assert(split_dim_size % num_split == 0); in SplitOpKernel()
85 int32 piece_size = split_dim_size / num_split; in SplitOpKernel()
173 const int32 num_split = output_ptr_data.size; in SplitVOpKernel_fixed() local
180 int32 piece_size = suffix_dim_size / num_split; in SplitVOpKernel_fixed()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
H A Dsplit_op.cc38 const int32_t num_split = num_outputs(); in Compile() local
58 ctx, num_split > 0, in Compile()
60 "Number of ways to split should be > 0, but got ", num_split)); in Compile()
63 ctx, input_shape.dim_size(split_dim) % num_split == 0, in Compile()
68 "and num_split ", num_split)); in Compile()
72 const int32_t slice_size = input_shape.dim_size(split_dim) / num_split; in Compile()
89 for (int i = 0; i < num_split; ++i) { in Compile()
105 const int32_t num_split = num_outputs(); in Compile() local
129 ctx, num_split > 0, in Compile()
131 "Number of ways to split should be > 0, but got ", num_split)); in Compile()
[all …]
/aosp_15_r20/external/tensorflow/tensorflow/python/kernel_tests/sparse_ops/
H A Dsparse_split_op_test.py85 sp_input=self._SparseTensor_4x6(), num_split=2, axis=axis))
102 sp_input=self._SparseTensor_5x7(), num_split=3, axis=axis))
118 sp_input=self._SparseTensor_5x7(), num_split=4, axis=axis)
140 sp_input=self._SparseTensor_5x7(), num_split=2, axis=axis))
154 sp_input=self._SparseTensor_5x7(), num_split=3, axis=axis)
172 sp_input=self._SparseTensor_4x6(), num_split=4, axis=axis))
193 sp_input=self._SparseTensor_4x6(), num_split=3, axis=axis))
212 sp_input=self._SparseTensor_4x6(), num_split=6, axis=axis))
238 sp_input=sp_input, num_split=2, axis=axis)
250 sp_input=self._SparseTensor_4x6(), num_split=3, axis=axis))
[all …]
/aosp_15_r20/external/tensorflow/tensorflow/core/util/sparse/
H A Dsparse_tensor.h174 const int num_split, std::vector<SparseTensor>* result);
497 const int split_dim, const int num_split, in Split() argument
502 output_indices.reserve(num_split); in Split()
503 output_values.reserve(num_split); in Split()
504 output_shapes.reserve(num_split); in Split()
508 output_indices_t.reserve(num_split); in Split()
509 output_values_t.reserve(num_split); in Split()
513 std::vector<int> num_values(num_split, 0); in Split()
516 const int split_size = split_dim_size / num_split; in Split()
518 if (!(num_split > 0 && num_split <= split_dim_size)) { in Split()
[all …]
/aosp_15_r20/external/tensorflow/tensorflow/python/kernel_tests/array_ops/
H A Dsplit_op_test.py129 num_split = np.random.randint(16, 25)
131 num_split = np.random.randint(2, 8)
132 size_splits = np.random.randint(2, 8, num_split, dtype=np.int32)
139 for i in range(num_split):
156 num_split = 1000
157 size_splits = np.random.randint(1, 3, num_split, dtype=np.int32)
165 for i in range(num_split):
287 num_split = np.random.randint(9, 15)
289 num_split = np.random.randint(2, 8)
290 shape[split_dim] = np.random.randint(2, 5) * num_split
[all …]
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tfr/examples/pad/
H A Dops_defs.py63 num_split=2)
68 num_split=2)
74 num_split=2)
79 num_split=2)
110 num_split=3)
122 num_split=2)
127 num_split=2)
/aosp_15_r20/external/tensorflow/tensorflow/core/api_def/base_api/
H A Dapi_def_SparseSplit.pbtxt45 name: "num_split"
50 summary: "Split a `SparseTensor` into `num_split` tensors along one dimension."
52 If the `shape[split_dim]` is not an integer multiple of `num_split`. Slices
53 `[0 : shape[split_dim] % num_split]` gets one extra dimension.
54 For example, if `split_dim = 1` and `num_split = 2` and the input is
H A Dapi_def_Split.pbtxt22 `values.shape[split_dim] / num_split`.
26 name: "num_split"
32 summary: "Splits a tensor into `num_split` tensors along one dimension."
/aosp_15_r20/external/tensorflow/tensorflow/core/ops/compat/ops_history_v2/
H A DSparseSplit.pbtxt22 number_attr: "num_split"
27 number_attr: "num_split"
32 number_attr: "num_split"
35 name: "num_split"
H A DSplit.pbtxt14 number_attr: "num_split"
17 name: "num_split"
H A DSplitV.pbtxt18 number_attr: "num_split"
21 name: "num_split"
/aosp_15_r20/external/tensorflow/tensorflow/core/ops/compat/ops_history_v1/
H A DSparseSplit.pbtxt22 number_attr: "num_split"
27 number_attr: "num_split"
32 number_attr: "num_split"
35 name: "num_split"
H A DSplit.pbtxt14 number_attr: "num_split"
17 name: "num_split"
H A DSplitV.pbtxt18 number_attr: "num_split"
21 name: "num_split"
/aosp_15_r20/external/tensorflow/tensorflow/java/src/test/java/org/tensorflow/
H A DGraphOperationTest.java182 private static int split(int[] values, int num_split) { in split() argument
187 .setAttr("num_split", num_split) in split()
193 private static int splitWithInputList(int[] values, int num_split, String name) { in splitWithInputList() argument
198 .setAttr("num_split", num_split) in splitWithInputList()
/aosp_15_r20/external/mesa3d/src/gallium/auxiliary/gallivm/
H A Dlp_bld_pack.c650 int num_split = src_type.width * src_type.length / 128; in lp_build_pack2() local
662 assert(num_split <= LP_MAX_VECTOR_WIDTH / 128); in lp_build_pack2()
664 for (i = 0; i < num_split / 2; i++) { in lp_build_pack2()
675 for (i = 0; i < num_split / 2; i++) { in lp_build_pack2()
680 tmpres[i+num_split/2] = lp_build_intrinsic_binary(builder, intrinsic, in lp_build_pack2()
684 tmpres[i+num_split/2] = LLVMBuildBitCast(builder, tmpres[i+num_split/2], in lp_build_pack2()
688 res = lp_build_concat(gallivm, tmpres, ndst_type, num_split); in lp_build_pack2()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tfr/python/
H A Dtfr_gen_test.py75 z, _ = array_ops.Split(axis=0, value=x, num_split=2)
82 z = array_ops.Split(axis=0, value=x, num_split=2)
139 value=lhs, size_splits=[rhs, -1], axis=0, num_split=2)
141 value=lhs, size_splits=[rhs, rhs], axis=1, num_split=2)
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/
H A Dxla_sharding_util.cc45 mlir::LogicalResult CreateSplitOp(const int num_split, in CreateSplitOp() argument
70 if (shape[split_dimension] % num_split != 0) { in CreateSplitOp()
76 split_dimension, num_split)); in CreateSplitOp()
79 shape[split_dimension] = shape[split_dimension] / num_split; in CreateSplitOp()
88 llvm::SmallVector<mlir::Type, 4> output_types(num_split, output_type); in CreateSplitOp()
93 builder->getIntegerAttr(builder->getIntegerType(32), num_split)); in CreateSplitOp()
/aosp_15_r20/external/tensorflow/tensorflow/core/transforms/constant_folding/tests/
H A Dsplit_removal.mlir11 …%Split, %ctl_3 = Split(%Const, %VariableV2) name("s1") {T = f32, num_split = 1 : i64} : (tensor<i3…
12 …%Split_4:2, %ctl_5 = Split(%Const, %VariableV2_0) name("s2") {T = f32, num_split = 2 : i64} : (ten…

1234