| /external/tensorflow/tensorflow/compiler/xla/service/ |
| hlo_memory_scheduler_test.cc | 52 const Shape vec = ShapeUtil::MakeShape(xla::F32, {42});
72 return ShapeUtil::ByteSizeOf(buffer.shape());
123 return ShapeUtil::ByteSizeOf(buffer.shape(), /*pointer_size=*/8);
153 const Shape r1f32 = ShapeUtil::MakeShape(xla::F32, {6});
181 return ShapeUtil::ByteSizeOf(buffer.shape(),
197 const Shape r1f32 = ShapeUtil::MakeShape(xla::F32, {5});
232 return ShapeUtil::ByteSizeOf(buffer.shape(), 2);
247 const Shape r1f32 = ShapeUtil::MakeShape(F32, {4});
258 HloInstruction::CreateCompare(ShapeUtil::MakeShape(PRED, {}), cond_param,
285 return ShapeUtil::ByteSizeOf(buffer.shape())
[all...] |
| instruction_fusion_test.cc | 116 HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(S32, {}), "0"));
118 HloInstruction::CreateReshape(ShapeUtil::MakeShape(S32, {1, 1}), param0));
132 HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(S32, {}), "0"));
134 HloInstruction::CreateReshape(ShapeUtil::MakeShape(S32, {1, 1}), param0));
148 HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(S32, {}), "0"));
150 ShapeUtil::MakeShape(S32, {}), param0, {}));
163 auto shape = ShapeUtil::MakeShape(F32, {16, 16});
354 auto shape = ShapeUtil::MakeShape(F32, {16, 16});
374 auto shape = ShapeUtil::MakeShape(F32, {16, 16});
375 auto small_shape = ShapeUtil::MakeShape(F32, {16})
[all...] |
| hlo_computation_test.cc | 71 Shape r0f32_ = ShapeUtil::MakeShape(F32, {});
503 Shape r0s64 = ShapeUtil::MakeShape(S64, {});
504 Shape r0s32 = ShapeUtil::MakeShape(S32, {});
505 Shape r0u32 = ShapeUtil::MakeShape(U32, {});
513 HloInstruction::CreateCompare(ShapeUtil::MakeShape(PRED, {}), param0,
528 ShapeUtil::Equal(clone->parameter_instruction(0)->shape(), r0f32_));
530 ShapeUtil::Equal(clone->parameter_instruction(1)->shape(), r0f32_));
532 ShapeUtil::Equal(clone->parameter_instruction(2)->shape(), r0s32));
534 ShapeUtil::Equal(clone->parameter_instruction(3)->shape(), r0u32));
538 const Shape s1 = ShapeUtil::MakeShape(F32, {5, 10})
[all...] |
| while_loop_invariant_code_motion_test.cc | 62 auto scalar_s32 = ShapeUtil::MakeShape(S32, {});
64 ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32, scalar_s32});
104 auto scalar_s32 = ShapeUtil::MakeShape(S32, {});
106 ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32, scalar_s32});
177 auto scalar_s32 = ShapeUtil::MakeShape(S32, {});
178 Shape while_shape = ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32});
215 auto scalar_s32 = ShapeUtil::MakeShape(S32, {});
217 ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32, scalar_s32});
253 auto scalar_s32 = ShapeUtil::MakeShape(S32, {});
254 auto token_shape = ShapeUtil::MakeTokenShape()
[all...] |
| algebraic_simplifier.cc | 123 return ShapeUtil::TransposeIsBitcast(operand->shape(), transpose->shape(),
269 ShapeUtil::DeleteDimension(dim, hlo->shape()), hlo));
278 ShapeUtil::MakeShape(hlo->shape().element_type(),
279 {ShapeUtil::ElementsIn(hlo->shape())}),
291 ShapeUtil::ChangeElementType(hlo->shape(), element_type), hlo));
319 Shape shape = ShapeUtil::FilterDimensions(
345 // (ShapeUtil::Equal). If options_.is_layout_sensitive() is false, then the
346 // tests shape compatibility (ShapeUtil::Compatible).
405 Shape shape = ShapeUtil::MakeShape(F32, {});
457 return ShapeUtil::Equal(lhs->shape(), rhs->shape())
[all...] |
| hlo_creation_utils.cc | 116 Shape new_shape = ShapeUtil::MakeShape(operand->shape().element_type(),
131 ShapeUtil::MakeShape(start_indices->shape().element_type(), {1}),
135 ShapeUtil::MakeShape(start_indices->shape().element_type(), {}),
139 rank, ShapeUtil::MakeShape(start_indices->shape().element_type(), {}));
159 ShapeUtil::MakeShape(start_indices->shape().element_type(), {1}),
163 ShapeUtil::MakeShape(start_indices->shape().element_type(), {}),
167 rank, ShapeUtil::MakeShape(start_indices->shape().element_type(), {}));
180 Shape broadcast_shape = ShapeUtil::MakeShape(operand->shape().element_type(),
258 auto scalar_shape = ShapeUtil::MakeShape(operand->shape().element_type(), {});
330 ShapeUtil::MakeShape(operand_shape.element_type(), new_shape_dims)
[all...] |
| hlo_sharding.cc | 61 CHECK(tuple_shape.IsTuple()) << ShapeUtil::HumanString(tuple_shape);
74 CHECK(tuple_shape.IsTuple()) << ShapeUtil::HumanString(tuple_shape);
202 // Empty tuples have no leaf nodes as far as ShapeUtil and ShapeTree are
205 return ShapeUtil::IsEmptyTuple(shape) ? 1 : ShapeUtil::GetLeafCount(shape);
211 << "Shape " << ShapeUtil::HumanString(shape) << " has " << shape_leaves
225 if (ShapeUtil::IsEmptyTuple(shape)) {
283 ShapeUtil::GetSubshape(shape, index_to_sharding.first), num_devices);
301 " against shape ", ShapeUtil::HumanString(shape)));
347 ToString(), ", input_shape=", ShapeUtil::HumanString(shape))
[all...] |
| hlo_verifier_test.cc | 70 const Shape scalar_shape = ShapeUtil::MakeShape(F32, {});
89 const Shape scalar_shape = ShapeUtil::MakeShape(F32, {});
108 const Shape scalar_shape = ShapeUtil::MakeShape(F32, {});
132 Shape s1 = ShapeUtil::MakeShape(F32, {1});
133 Shape s2 = ShapeUtil::MakeShape(F32, {2});
303 0, ShapeUtil::MakeShape(F32, {100}), "param"));
307 ShapeUtil::MakeShape(F32, {100}), param,
327 0, ShapeUtil::MakeShape(F32, {100}), "param"));
331 ShapeUtil::MakeShape(F32, {100}), param,
|
| instruction_fusion.cc | 51 ShapeUtil::ByteSizeOf(instruction.operand(0)->shape()) <
52 ShapeUtil::ByteSizeOf(instruction.shape());
113 return ShapeUtil::ElementIsComplex(instruction.shape());
174 ShapeUtil::ForEachSubshape(
178 output_rank = std::max(output_rank, ShapeUtil::TrueRank(subshape));
188 ShapeUtil::IsEffectiveScalar(operand->shape())) {
191 return ShapeUtil::TrueRank(operand->shape()) >= output_rank;
273 ShapeUtil::ForEachSubshape(
277 size += ShapeUtil::ElementsIn(subshape);
|
| bfloat16_propagation.cc | 52 ShapeUtil::ForEachSubshape(
100 ShapeUtil::ForEachSubshape(
176 ShapeUtil::ForEachSubshape(
209 const Shape& subshape = ShapeUtil::GetSubshape(hlo.shape(), index);
341 ShapeUtil::ForEachSubshape(hlo->shape(), [&](const Shape& /* subshape */,
376 ShapeUtil::ForEachSubshape(
415 ShapeUtil::ForEachSubshape(
419 if (!ShapeUtil::IsLeafIndex(parameter->shape(), index)) {
456 ShapeUtil::ForEachSubshape(root->shape(), [this, hlo, root, output](
459 if (!ShapeUtil::IsLeafIndex(hlo->shape(), index))
[all...] |
| buffer_liveness_test.cc | 72 ShapeUtil::Compatible(ShapeUtil::GetSubshape(a->shape(), index),
73 ShapeUtil::GetSubshape(b->shape(), index)));
102 const Shape vec_ = ShapeUtil::MakeShape(xla::F32, {42});
484 ShapeUtil::MakeTupleShape(
485 {ShapeUtil::MakeShape(F32, {8}), ShapeUtil::MakeShape(S32, {4})}),
492 ShapeUtil::GetSubshape(tuple_param0->shape(), {0});
504 ShapeUtil::GetSubshape(tuple_param0->shape(), {1});
548 ShapeUtil::MakeTupleShape
[all...] |
| hlo_cost_analysis.cc | 50 // sizes of the inputs and outputs. The default ShapeUtil::ByteSizeOf does not
90 auto computation_count = ShapeUtil::ElementsIn(shape);
251 kFmaFlops * ShapeUtil::ElementsIn(dot_shape) * reduction_width;
269 const int64 element_count = ShapeUtil::ElementsIn(map->shape());
293 ShapeUtil::ElementsIn(arg->shape()) - ShapeUtil::ElementsIn(output_shape);
317 ShapeUtil::ElementsIn(reduce_window->shape());
341 const auto source_element_count = ShapeUtil::ElementsIn(source->shape());
448 ShapeUtil::GetDimension(lhs_shape, input_feature_dim);
450 ShapeUtil::GetDimension(result_shape, output_feature_dim)
[all...] |
| hlo_subcomputation_unification_test.cc | 61 Shape r0s32_ = ShapeUtil::MakeShape(S32, {});
62 Shape r0f32_ = ShapeUtil::MakeShape(S32, {});
63 Shape r1s32_5_ = ShapeUtil::MakeShape(S32, {5});
64 Shape r1s32_3_ = ShapeUtil::MakeShape(S32, {3});
142 ShapeUtil::MakeShape(S32, {8}), {x, y}, 0));
|
| /external/tensorflow/tensorflow/compiler/xla/tests/ |
| multioutput_fusion_test.cc | 67 const Shape elem_shape0 = ShapeUtil::MakeShapeWithLayout(F32, {}, {});
69 ShapeUtil::MakeShapeWithLayout(F32, {size, size}, {1, 0});
112 Literal arg1(ShapeUtil::MakeShapeWithDescendingLayout(F32, {size, size}));
115 Literal expect(ShapeUtil::MakeShapeWithDescendingLayout(F32, {size, size}));
128 ShapeUtil::MakeShapeWithDescendingLayout(F32, {size});
130 ShapeUtil::MakeShapeWithDescendingLayout(F64, {size});
150 ShapeUtil::MakeShapeWithDescendingLayout(F32, {size, 1}), add));
155 ShapeUtil::MakeShapeWithDescendingLayout(F32, {1}), sub, reshape,
176 Literal input0(ShapeUtil::MakeShapeWithDescendingLayout(F32, {size}));
178 Literal input1(ShapeUtil::MakeShapeWithDescendingLayout(F64, {size}))
[all...] |
| conditional_test.cc | 163 Shape r0f32_ = ShapeUtil::MakeShape(F32, {});
164 Shape r1s2f32_ = ShapeUtil::MakeShape(F32, {2});
165 Shape tuple_2_r0f32_ = ShapeUtil::MakeTupleShape(
166 {ShapeUtil::MakeShape(F32, {}), ShapeUtil::MakeShape(F32, {})});
167 Shape tuple_2_r1s2f32_ = ShapeUtil::MakeTupleShape(
168 {ShapeUtil::MakeShape(F32, {2}), ShapeUtil::MakeShape(F32, {2})});
169 Shape empty_tuple_ = ShapeUtil::MakeTupleShape({});
343 Shape r0bool = ShapeUtil::MakeShape(PRED, {})
[all...] |
| map_test.cc | 55 auto x = Parameter(&mapped_builder, 0, ShapeUtil::MakeShape(F32, {}), "x");
65 auto lhs = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {}), "x");
66 auto rhs = Parameter(&b, 1, ShapeUtil::MakeShape(F32, {}), "y");
78 (void)Parameter(&mapped_builder, 0, ShapeUtil::MakeShape(F32, {}), "x");
92 auto x = Parameter(&mapped_builder, 0, ShapeUtil::MakeShape(F32, {}), "x");
110 auto x = Parameter(&mapped_builder, 0, ShapeUtil::MakeShape(F32, {}), "x");
128 auto x = Parameter(&builder, 0, ShapeUtil::MakeShape(F32, {}), "x");
141 auto x = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {}), "x");
142 auto y = Parameter(&b, 1, ShapeUtil::MakeShape(F32, {}), "y");
158 auto x = Parameter(&mapped_builder, 0, ShapeUtil::MakeShape(F32, {}), "x")
[all...] |
| tuple_test.cc | 179 ASSERT_TRUE(ShapeUtil::Equal(vector_shape, ShapeUtil::MakeShape(F32, {3})));
180 ASSERT_TRUE(ShapeUtil::Equal(matrix_shape,
181 ShapeUtil::MakeShape(F32, {/*y=*/2, /*x=*/3})));
303 auto x = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {}), "x");
433 Shape data_shape = ShapeUtil::MakeShape(F32, {3});
434 Shape inner_tuple_shape = ShapeUtil::MakeTupleShape({data_shape, data_shape});
436 ShapeUtil::MakeTupleShape({inner_tuple_shape, data_shape});
462 Shape c64r0 = ShapeUtil::MakeShape(C64, {});
463 Shape c64r1 = ShapeUtil::MakeShape(C64, {2})
[all...] |
| token_hlo_test.cc | 75 HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(F32, {}), "p0"));
77 HloInstruction::CreateParameter(1, ShapeUtil::MakeTokenShape(), "p1"));
97 ShapeUtil::MakeTupleShape(
98 {ShapeUtil::MakeShape(F32, {1, 2, 3}), ShapeUtil::MakeTokenShape()}),
|
| /external/tensorflow/tensorflow/compiler/xrt/tests/ |
| raw_api_test.cc | 161 xla::ShapeUtil::MakeShape(xla::F32, {2}), "P0");
163 xla::ShapeUtil::MakeShape(xla::F32, {2}), "P1");
174 xla::ShapeUtil::MakeShapeWithLayout(xla::F32, {2, 2}, {0, 1}), "P0");
177 xla::ShapeUtil::MakeShapeWithLayout(xla::F32, {2, 1}, {0, 1}), "P1");
187 auto p0 = xla::Parameter(&builder, 0, xla::ShapeUtil::MakeShape(xla::S64, {}),
189 auto p1 = xla::Parameter(&builder, 1, xla::ShapeUtil::MakeShape(xla::S64, {}),
198 xla::ShapeUtil::MakeShape(xla::F32, {2}), "P0");
200 xla::ShapeUtil::MakeShape(xla::F32, {2}), "P1");
208 auto p0 = xla::Parameter(&builder, 0, xla::ShapeUtil::MakeShape(xla::F32, {}),
210 auto p1 = xla::Parameter(&builder, 1, xla::ShapeUtil::MakeShape(xla::F32, {})
[all...] |
| /external/tensorflow/tensorflow/compiler/xla/ |
| literal_util.cc | 50 ShapeUtil::ForEachMutableSubshape(
62 ShapeUtil::ForEachSubshape(
88 ShapeUtil::MakeShape(primitive_type, dimensions));
102 return Literal(ShapeUtil::MakeTokenShape());
273 ShapeUtil::MakeShape(PRED, {static_cast<int64>(values.bits())}));
279 Literal literal(ShapeUtil::MakeShape(U8, {static_cast<int64>(value.size())}));
299 CHECK_EQ(ShapeUtil::ElementsIn(literal.shape()), new_num_elements);
303 ShapeUtil::MakeShape(literal.shape().element_type(), new_dimensions));
312 for (int64 i = 0; i < ShapeUtil::ElementsIn(literal.shape()); ++i) {
370 CHECK_GT(ShapeUtil::ElementsIn(literal.shape()), 0)
[all...] |
| /external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| qr.cc | 164 const int64 m = ShapeUtil::GetDimension(a_shape, -2);
165 const int64 n = ShapeUtil::GetDimension(a_shape, -1);
170 batch_dims[i] = ShapeUtil::GetDimension(a_shape, i);
222 builder, ShapeUtil::MakeShape(type, ConcatVectors(batch_dims, {m, n})));
224 ShapeUtil::MakeShape(type, ConcatVectors(batch_dims, {n})));
290 ShapeUtil::MakeShape(type, ConcatVectors(batch_dims, {m, n})));
335 const int64 m = ShapeUtil::GetDimension(a_shape, -2);
336 const int64 n = ShapeUtil::GetDimension(a_shape, -1);
347 batch_dims[i] = ShapeUtil::GetDimension(a_shape, i);
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| instruction_fusion.cc | 34 !ShapeUtil::IsScalar(constant->shape())) {
48 return !ShapeUtil::ElementIsFloating(instruction.shape()) &&
91 int64 num_output_buffers = ShapeUtil::SubshapeCount(a->shape()) +
92 ShapeUtil::SubshapeCount(b->shape());
246 return ShapeUtil::IsEffectiveScalar(producer->shape()) &&
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| reduce_window_op.cc | 97 xla::ShapeUtil::Compatible(reducer.xla_output_shape, scalar_shape),
100 xla::ShapeUtil::HumanString(scalar_shape), " got ",
101 xla::ShapeUtil::HumanString(reducer.xla_output_shape)));
|
| /external/tensorflow/tensorflow/compiler/tf2xla/ |
| xla_helpers.cc | 70 xla::ShapeUtil::MakeShape(input.shape().element_type(), dimensions);
71 int64 elements_before = xla::ShapeUtil::ElementsIn(input.shape());
72 int64 elements_after = xla::ShapeUtil::ElementsIn(shape);
|
| /external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| conv_canonicalization.cc | 73 ShapeUtil::MakeShape(input->shape().element_type(), new_input_dims);
95 ShapeUtil::MakeShape(kernel->shape().element_type(), new_kernel_dims);
114 ShapeUtil::MakeShape(hlo->shape().element_type(), new_conv_dims);
|
