| /external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| sorting.cc | 30 Shape iota_shape = local
32 XlaOp iota_s32 = Iota(builder, iota_shape, last_dim);
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| matrix_band_part_op.cc | 64 xla::Shape iota_shape = xla::ShapeUtil::MakeShape(index_xla_type, {m, n}); variable
65 xla::XlaOp iota_m = xla::Iota(builder, iota_shape, /*iota_dimension=*/0);
66 xla::XlaOp iota_n = xla::Iota(builder, iota_shape, /*iota_dimension=*/1);
|
| conv_op_helpers.cc | 108 xla::Shape iota_shape = xla::ShapeUtil::MakeShape(xla::S32, iota_dimensions); local
110 builder, iota_shape, /*iota_dimension=*/iota_dimensions.size() - 2);
112 builder, iota_shape, /*iota_dimension=*/iota_dimensions.size() - 1);
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| stable_sort_expander.cc | 64 Shape iota_shape = sort->operand(0)->shape(); local
69 if (iota_shape.dimensions(sort->sort_dimension()) >
74 iota_shape.set_element_type(S32);
76 HloInstruction::CreateIota(iota_shape, sort->sort_dimension()));
104 new_shapes.push_back(iota_shape);
|
| /external/tensorflow/tensorflow/compiler/tf2xla/ |
| xla_helpers.cc | 96 xla::Shape iota_shape; local
98 TensorShapeToXLAShape(index_type, output_shape, &iota_shape));
102 xla::Eq(indices, xla::Iota(builder, iota_shape, axis), broadcast_dims),
|
