Home
last modified time | relevance | path

Searched refs:indices_shape (Results 1 – 25 of 52) sorted by relevance

123

/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dgather_op.cc30 const xla::XlaOp& indices, const TensorShape& indices_shape, in XlaGather() argument
49 CHECK_GE(indices_shape.dims(), 1); in XlaGather()
50 num_index_dims = indices_shape.dim_size(indices_shape.dims() - 1); in XlaGather()
51 for (int64 i = 0, e = indices_shape.dims() - 1; i < e; i++) { in XlaGather()
52 num_indices *= indices_shape.dim_size(i); in XlaGather()
56 for (int64 i = 0, e = indices_shape.dims(); i < e; i++) { in XlaGather()
57 num_indices *= indices_shape.dim_size(i); in XlaGather()
68 TensorShape indices_shape_no_index_vectors{indices_shape}; in XlaGather()
134 indices_are_nd ? (indices_shape.dims() - 1) : indices_shape.dims(); in XlaGather()
139 dim_numbers.set_index_vector_dim(indices_are_nd ? (indices_shape.dims() - 1) in XlaGather()
[all …]
Dscatter_nd_op.cc33 const TensorShape& indices_shape, in ValidateUpdateShape() argument
35 if (indices_shape.dims() < 1) { in ValidateUpdateShape()
38 indices_shape.DebugString()); in ValidateUpdateShape()
41 const int64 num_index_dims = indices_shape.dim_size(indices_shape.dims() - 1); in ValidateUpdateShape()
42 const int64 batch_dim = indices_shape.dims() - 1; in ValidateUpdateShape()
49 ", indices.shape: ", indices_shape.DebugString(), in ValidateUpdateShape()
64 if (updates_shape.dim_size(d) != indices_shape.dim_size(d)) { in ValidateUpdateShape()
84 TensorShape indices_shape = context->InputShape(0); in Compile() local
97 buffer_shape.num_elements() > 0 || (indices_shape.num_elements() == 0 && in Compile()
101 indices_shape.DebugString())); in Compile()
[all …]
Dsparse_to_dense_op.cc31 const TensorShape indices_shape = context->InputShape(0); in Compile() local
32 OP_REQUIRES(context, indices_shape.dims() <= 2, in Compile()
36 indices_shape.DebugString())); in Compile()
38 indices_shape.dims() > 0 ? indices_shape.dim_size(0) : 1; in Compile()
40 indices_shape.dims() > 1 ? indices_shape.dim_size(1) : 1; in Compile()
Ddynamic_stitch_op.cc71 TensorShape indices_shape; in Compile() local
74 &indices_shape)); in Compile()
77 ctx, TensorShapeUtils::StartsWith(data_shape, indices_shape), in Compile()
81 "].shape = ", indices_shape.DebugString())); in Compile()
85 data_shape, indices_shape), in Compile()
88 input_num, "].shape[", indices_shape.dims(), in Compile()
93 "].shape = ", indices_shape.DebugString())); in Compile()
97 {indices_shape.num_elements()}, in Compile()
Dsegment_reduction_ops.cc57 TensorShape indices_shape = ctx->InputShape(1); in Compile() local
62 OP_REQUIRES(ctx, data_shape.dims() >= indices_shape.dims(), in Compile()
67 for (int d = 0; d < indices_shape.dims(); ++d) { in Compile()
69 ctx, (data_shape.dim_size(d) == indices_shape.dim_size(d)), in Compile()
74 " vs. ", indices_shape.dim_size(d))); in Compile()
78 buffer_shape.RemoveDimRange(0, indices_shape.dims()); in Compile()
Done_hot_op.cc33 const TensorShape indices_shape = ctx->InputShape(0); in Compile() local
38 const int indices_dims = indices_shape.dims(); in Compile()
68 indices_shape, ctx->Input(0), ctx->Input(2), in Compile()
Dtensor_array_ops.cc305 const TensorShape indices_shape = ctx->InputShape(1); in Compile() local
306 OP_REQUIRES(ctx, indices_shape.dims() == 1, in Compile()
342 XlaGather(ta, ta_shape, indices, indices_shape, /*axis=*/0, in Compile()
373 const TensorShape indices_shape = ctx->InputShape(1); in Compile() local
374 OP_REQUIRES(ctx, indices_shape.dims() >= 1, in Compile()
376 const int num_indices = indices_shape.dim_size(0); in Compile()
/external/tensorflow/tensorflow/core/kernels/
Dsparse_add_op_test.cc63 const auto indices_shape = TensorShape({4, 2}); in TEST_F() local
70 AddInputFromArray<int64>(indices_shape, indices); \ in TEST_F()
81 Tensor expected_indices(allocator(), DT_INT64, indices_shape); in TEST_F()
103 const auto indices_shape = TensorShape({4, 2}); \
109 AddInputFromArray<int64>(indices_shape, indices); \
160 const auto indices_shape = TensorShape({4, 2}); \
166 auto AddSparseTensor = [indices, indices_shape, shape, \
168 AddInputFromArray<int64>(indices_shape, indices); \
Done_hot_op.cc55 const TensorShape& indices_shape = indices.shape(); in Compute() local
57 const int indices_dims = indices_shape.dims(); in Compute()
84 MultiplyWithoutOverflow(indices_shape.num_elements(), depth_v) >= 0, in Compute()
86 indices_shape.DebugString(), " + [", depth_v, in Compute()
89 TensorShape output_shape = indices_shape; in Compute()
104 prefix_dim_size *= indices_shape.dim_size(i); in Compute()
106 TI suffix_dim_size = indices_shape.num_elements() / prefix_dim_size; in Compute()
Dgather_nd_op.cc96 const TensorShape& indices_shape(indices.shape()); in DoGatherNd() local
97 const int64 indices_nd = indices_shape.dim_size(indices_shape.dims() - 1); in DoGatherNd()
101 for (int i = 0; i < indices_shape.dims() - 1; ++i) { in DoGatherNd()
102 N_big *= indices_shape.dim_size(i); in DoGatherNd()
119 for (int i = 0; i < indices_shape.dims() - 1; ++i) { in DoGatherNd()
120 N_result *= indices_shape.dim_size(i); in DoGatherNd()
126 TensorShape result_shape(indices_shape); in DoGatherNd()
Dsparse_dense_binary_op_shared_test.cc99 const auto indices_shape = TensorShape({4, 2}); in TEST_F() local
110 AddInputFromArray<int64>(indices_shape, indices); in TEST_F()
128 const auto indices_shape = TensorShape({4, 2}); in TEST_F() local
138 AddInputFromArray<int64>(indices_shape, indices); in TEST_F()
155 const auto indices_shape = TensorShape({4, 2}); in TEST_F() local
165 AddInputFromArray<int64>(indices_shape, indices); in TEST_F()
187 const auto indices_shape = TensorShape({4, 2}); in TEST_F() local
198 AddInputFromArray<int64>(indices_shape, indices); in TEST_F()
Dsparse_reduce_sum_op_test.cc53 const auto indices_shape = TensorShape({4, 2}); in TEST_F() local
59 AddInputFromArray<int64>(indices_shape, indices); in TEST_F()
95 const auto indices_shape = TensorShape({4, 2}); in TEST_F() local
101 AddInputFromArray<int64>(indices_shape, indices); in TEST_F()
Dragged_gather_op_test.cc33 const TensorShape& indices_shape, const std::vector<INDEX_TYPE>& indices, in BuildRaggedGatherGraph() argument
40 int64 num_splits = PARAMS_RAGGED_RANK + indices_shape.dims() - 1; in BuildRaggedGatherGraph()
58 AddInputFromArray<INDEX_TYPE>(indices_shape, indices); in BuildRaggedGatherGraph()
Dscatter_nd_op.cc558 const TensorShape& indices_shape(indices.shape()); in PrepareAndValidateInputs() local
567 indices_shape.num_elements(), in PrepareAndValidateInputs()
577 "must match. Got indices.shape ", indices_shape.DebugString(), in PrepareAndValidateInputs()
598 *slice_dim = (indices_shape.dims() > 1) in PrepareAndValidateInputs()
599 ? indices_shape.dim_size(indices_shape.dims() - 1) in PrepareAndValidateInputs()
621 *num_updates = indices_shape.num_elements() / safe_slice_dim; in PrepareAndValidateInputs()
/external/tensorflow/tensorflow/compiler/tf2xla/lib/
Dscatter.cc41 TF_ASSIGN_OR_RETURN(xla::Shape indices_shape, builder->GetShape(indices)); in XlaScatter()
43 xla::AsInt64Slice(indices_shape.dimensions()); in XlaScatter()
54 xla::ShapeUtil::HumanString(indices_shape), in XlaScatter()
140 ? indices_shape.dimensions_size() - 1 in XlaScatter()
141 : indices_shape.dimensions_size()); in XlaScatter()
190 VLOG(3) << " Indices: " << xla::ShapeUtil::HumanString(indices_shape); in XlaScatter()
/external/tensorflow/tensorflow/python/kernel_tests/
Dbatch_scatter_ops_test.py58 for indices_shape in (2,), (3, 7), (3, 4, 7):
61 sparse_dim = len(indices_shape) - 1
63 indices_shape[sparse_dim], size=indices_shape, dtype=itype)
65 np.random.randn(*(indices_shape + extra_shape)), vtype)
67 old = _AsType(np.random.randn(*(indices_shape + extra_shape)), vtype)
Dgather_op_test.py97 for indices_shape in (), (0,), (2, 0), (2, 3):
101 indices = np.random.randint(shape[axis], size=indices_shape)
343 indices_shape=[2, 3, 8, 9, 10],
351 indices_shape=[2, 3, 8, 9, 10],
359 indices_shape=[2, 3, 8, 9, 10],
367 indices_shape=[2, 3, 8, 9, 10],
375 indices_shape=[2, 3, 8, 9, 10],
383 indices_shape=[2, 3, 8, 9, 10],
391 indices_shape=[2, 3, 8, 9, 10],
399 indices_shape=[2, 3, 8, 9, 10],
[all …]
/external/tensorflow/tensorflow/compiler/xla/tests/
Dtest_utils_test.cc280 const Shape& indices_shape = args[1].shape(); in XLA_TEST_F() local
282 ShapeUtil::Equal(indices_shape, ShapeUtil::MakeShape(S32, {10, 2}))) in XLA_TEST_F()
283 << ShapeUtil::HumanString(indices_shape); in XLA_TEST_F()
318 const Shape& indices_shape = args[1].shape(); in XLA_TEST_F() local
320 ShapeUtil::Equal(indices_shape, ShapeUtil::MakeShape(S32, {10, 2}))) in XLA_TEST_F()
321 << ShapeUtil::HumanString(indices_shape); in XLA_TEST_F()
/external/tensorflow/tensorflow/compiler/tf2xla/
Dxla_helpers.cc85 DataType index_type, const TensorShape& indices_shape, in OneHot() argument
90 std::vector<int64> broadcast_dims(indices_shape.dims()); in OneHot()
94 TensorShape output_shape = indices_shape; in OneHot()
/external/tensorflow/tensorflow/python/ops/ragged/
Dragged_gather_ops.py177 indices_shape = indices.shape
178 indices_ndims = indices_shape.ndims
189 index_size = tensor_shape.dimension_value(indices_shape[-1])
/external/tensorflow/tensorflow/core/ops/
Dresource_variable_ops.cc267 ShapeHandle indices_shape = c->input(1); in __anon4d581ac00402() local
280 c->WithRankAtLeast(indices_shape, batch_dims, &unused)); in __anon4d581ac00402()
292 c->Subshape(indices_shape, batch_dims, &indices_subshape)); in __anon4d581ac00402()
319 ShapeHandle indices_shape = c->input(1); in ResourceScatterUpdateShape() local
325 TF_RETURN_IF_ERROR(c->Concatenate(indices_shape, var_subshape, &concat)); in ResourceScatterUpdateShape()
/external/tensorflow/tensorflow/python/framework/
Dsparse_tensor.py130 indices_shape = indices.get_shape().with_rank(2)
135 indices_shape.dims[0].merge_with(values_shape.dims[0])
138 indices_shape.dims[1].merge_with(dense_shape_shape.dims[0])
/external/tensorflow/tensorflow/lite/kernels/
Dembedding_lookup_sparse_test.cc36 std::initializer_list<int> indices_shape, in EmbeddingLookupSparseOpModel() argument
48 BuildInterpreter({lookup_shape, indices_shape, dense_shape_shape, in EmbeddingLookupSparseOpModel()
/external/tensorflow/tensorflow/compiler/tests/
Dscatter_nd_op_test.py82 for ref_shape, indices_shape in zip(ref_shapes, indices_shapes):
83 num_updates = indices_shape[0]
84 ixdim = indices_shape[-1]
/external/tensorflow/tensorflow/core/util/
Dexample_proto_helper.cc310 output_shapes->indices_shape.AddDim(total_num_features); in GetSparseTensorShapes()
311 output_shapes->indices_shape.AddDim(2); in GetSparseTensorShapes()
378 const TensorShape& indices_shape = sparse_tensor_batch_shapes.indices_shape; in BatchExampleProtoToTensors() local
383 Tensor(allocator, DT_INT64, indices_shape); in BatchExampleProtoToTensors()

123