| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | stable_sort_expander_test.cc | 109 m::Sort(m::Parameter(0), m::Iota()), 0))); in TEST_F()
155 m::Sort(m::Parameter(0), m::Iota()), 0))); in TEST_F()
187 m::Sort(m::Parameter(0), m::Parameter(1), m::Iota()), 0), in TEST_F()
189 m::Sort(m::Parameter(0), m::Parameter(1), m::Iota()), 1)))); in TEST_F()
253 m::Sort(m::Parameter(0), m::Iota(), m::Iota()), 0))); in TEST_F()
290 m::Sort(m::Parameter(0), m::Iota(), m::Iota()), 0))); in TEST_F()
321 m::Sort(m::Parameter(0), m::Iota()), 0))); in TEST_F()
352 m::Sort(m::Parameter(0), m::Iota()), 0)))); in TEST_F()
|
| D | qr_expander.cc | 97 XlaOp iota = Iota(builder, S32, m); in House()
210 auto iota_mn = Iota( in QrBlock()
230 auto iota = Reshape(Iota(a.builder(), S32, m), {m, 1}); in QrBlock()
234 auto successor_mask = Gt(Iota(a.builder(), S32, m), j); in QrBlock()
254 Iota(builder, ShapeUtil::MakeShape(S32, ConcatVectors(batch_dims, {n})), in QrBlock()
|
| D | cholesky_expander.cc | 86 Iota(body_builder, ShapeUtil::MakeShape(S32, matrix_dims), ndims - 1); in CholeskyUnblocked()
88 Iota(body_builder, ShapeUtil::MakeShape(S32, matrix_dims), ndims - 2); in CholeskyUnblocked()
|
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | outfeed_receiver_test.cc | 125 XlaOp data = Iota(&builder, shape0, 0); in TEST()
158 XlaOp data0 = Iota(&builder0, shape0, 0); in TEST()
168 XlaOp data1 = Iota(&builder1, shape1, 0); in TEST()
203 XlaOp data0 = Iota(&builder, shape0, 0); in TEST()
211 XlaOp data1 = Iota(&builder, shape1, 0); in TEST()
246 XlaOp data0 = Iota(&builder, shape0, 0); in TEST()
253 XlaOp data1 = Iota(&builder, shape1, 0); in TEST()
279 XlaOp data0 = Iota(&builder, shape0, 0); in TEST()
306 XlaOp data = Iota(&builder, shape0, 0); in TEST()
|
| /external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| D | logdet.cc | 45 auto rows = Iota(a.builder(), ShapeUtil::ChangeElementType(a_shape, S32), in LogDet()
47 auto cols = Iota(a.builder(), ShapeUtil::ChangeElementType(a_shape, S32), in LogDet()
55 rows = Iota(a.builder(), row_shape, row_shape.rank() - 1); in LogDet()
|
| D | matrix.cc | 51 auto a = Iota(builder, U32, m); in IdentityMatrix()
52 auto b = Iota(builder, U32, n); in IdentityMatrix()
67 auto a = Iota(builder, S32, n); in GetDiagonalMask()
68 auto b = Iota(builder, S32, m) + ConstantR0WithType(builder, S32, diagonal); in GetDiagonalMask()
130 XlaOp diag_base_indices = BroadcastInDim(Iota(builder, S32, diag_len), in GetMatrixDiagonalViaGather()
221 auto a = Iota(builder, S32, n); in TriangleMask()
222 auto b = Iota(builder, S32, m) + ConstantR0<int32>(builder, diagonal); in TriangleMask()
276 mask = And(mask, Eq(Iota(builder, iota_shape, first_dim), in EinsumDiagonalMask()
277 Iota(builder, iota_shape, dim))); in EinsumDiagonalMask()
|
| D | slicing.cc | 191 Iota(builder, ShapeUtil::MakeShape(index_shape.element_type(), sizes), in TorchGather()
211 to_concat.push_back(Iota(builder, index_shape, i)); in TorchGather()
247 Iota(builder, in TorchScatterDense()
286 to_concat.push_back(Iota(builder, index_shape, batch_dim)); in TorchIndexSelect()
|
| D | sorting.cc | 50 XlaOp iota_s32 = Iota(builder, iota_shape, last_dim); in TopK()
92 XlaOp iota_s32 = Iota(builder, iota_shape, last_dim); in TopKWithPartitions()
|
| D | arithmetic.cc | 159 XlaOp iota = Iota(builder, iota_shape, axis); in ArgMinMax()
190 XlaOp iota = Iota( in ArgMinMaxTwoPass()
|
| D | quantize.h | 106 XlaOp iota_r1 = Iota(builder, U32, unpack_size);
|
| /external/tensorflow/tensorflow/compiler/xla/tests/ |
| D | iota_test.cc | 67 Iota(&builder, element_type, num_elements); in XLA_TEST_P()
104 Iota(&builder, ShapeUtil::MakeShape(element_type, dimensions), iota_dim); in XLA_TEST_P()
136 Iota(&builder, ShapeUtil::MakeShape(element_type, dimensions), iota_dim); in XLA_TEST_P()
|
| D | constants_test.cc | 186 auto val1 = Iota(&b, F32, 3); in TEST_F()
194 auto tuple = Tuple(&b, {Iota(&b, F32, 3), Iota(&b, F32, 1)}); in TEST_F()
|
| D | dynamism_inference_test.cc | 109 TEST_F(DynamismInferenceTest, Iota) { in TEST_F() argument
114 auto computation = Iota(&b, S32, 2); in TEST_F()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | sequence_ops.cc | 65 xla::Iota(builder, xla::primitive_util::NativeToPrimitiveType<T>(), in CreateRangeTensor()
162 xla::XlaOp iota = xla::Iota(ctx->builder(), ctx->output_xla_type(0), num); in Compile()
169 xla::XlaOp mask = xla::Iota(ctx->builder(), xla::S64, num); in Compile()
|
| D | matrix_band_part_op.cc | 65 xla::XlaOp iota_m = xla::Iota(builder, iota_shape, /*iota_dimension=*/0); in Compile()
66 xla::XlaOp iota_n = xla::Iota(builder, iota_shape, /*iota_dimension=*/1); in Compile()
|
| D | reverse_sequence_op.cc | 76 xla::XlaOp batch_idx = xla::Iota( in Compile()
80 xla::XlaOp forward_idx = xla::Iota( in Compile()
|
| D | extract_image_patches_op.cc | 121 xla::Eq(xla::Iota(builder, iota_kernel_shape, 0), in Compile()
122 xla::Iota(builder, iota_kernel_shape, 2)), in Compile()
|
| D | cast_op.cc | 164 auto iota = xla::Iota(ctx->builder(), xla_input_shape, in Compile()
190 xla::Iota(ctx->builder(), input_xla_shape, in Compile()
|
| D | where_op.cc | 64 xla::XlaOp iota = xla::Iota(ctx->builder(), iota_shape, axis); in Compile()
|
| D | in_topk_op.cc | 72 xla::Iota(xla_builder, targets_type_, predictions_shape.dim_size(1)); in Compile()
|
| D | diag_op.cc | 45 xla::XlaOp iota = xla::Iota(builder, xla::S32, last_dim_size); in CreateDiagonal()
|
| D | transpose_op.cc | 183 xla::Iota(ctx->builder(), in InvertPermutation()
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | topk_ops.cc | 152 xla::XlaOp iota_r1 = xla::Iota(builder, xla::S32, width); in CreateTopKUnique()
288 xla::XlaOp iota_r1 = xla::Iota(builder, xla::U32, width); in CreateMakeUnique()
|
| /external/tensorflow/tensorflow/core/kernels/image/ |
| D | non_max_suppression_op.cu.cc | 246 __global__ void Iota(const int num_elements, const T offset, T* to_fill) { in Iota() function
318 TF_CHECK_OK(GpuLaunchKernel(Iota<int>, config.block_count, in NmsGpu()
462 TF_CHECK_OK(GpuLaunchKernel(Iota<int>, config.block_count, in DoNMS()
|
| /external/tensorflow/tensorflow/core/kernels/ |
| D | training_ops_test.cc | 76 static Node* Iota(Graph* g, int n) { in Iota() function
173 auto indices = Iota(g, m); in SparseAdagrad()
|