Home
last modified time | relevance | path

Searched refs:return_shape (Results 1 – 9 of 9) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dir_emitter_nested.cc128 const Shape& return_shape = root_instruction->shape(); in CodegenNestedComputation() local
133 if (ShapeUtil::IsScalar(return_shape)) { in CodegenNestedComputation()
138 CHECK(return_shape.IsTuple()); in CodegenNestedComputation()
139 llvm::Type* tuple_type = llvm_ir::ShapeToIrType(return_shape, module_); in CodegenNestedComputation()
143 for (int i = 0; i < return_shape.tuple_shapes_size(); i++) { in CodegenNestedComputation()
144 const Shape& element_shape = return_shape.tuple_shapes(i); in CodegenNestedComputation()
Dir_emitter.cc638 const Shape& return_shape = computation.root_instruction()->shape(); in ComputeNestedElement() local
640 llvm_ir::ShapeToIrType(return_shape, module_), "return_buffer", &b_); in ComputeNestedElement()
649 if (!return_shape.IsTuple()) { in ComputeNestedElement()
653 llvm_ir::EmitTupleAllocasAtFunctionEntry(return_shape, &b_); in ComputeNestedElement()
654 llvm_ir::IrArray tuple_array(return_buffer, return_shape); in ComputeNestedElement()
/external/tensorflow/tensorflow/compiler/xla/
Dshape_util.cc910 const Shape* return_shape = &shape; in GetSubshape() local
912 CHECK(return_shape->IsTuple()) in GetSubshape()
914 return_shape = &return_shape->tuple_shapes(i); in GetSubshape()
916 return *return_shape; in GetSubshape()
921 const Shape* return_shape = &shape; in TryGetSubshape() local
923 if (!return_shape->IsTuple() || i < 0 || in TryGetSubshape()
924 i >= return_shape->tuple_shapes_size()) { in TryGetSubshape()
929 return_shape = &return_shape->tuple_shapes(i); in TryGetSubshape()
931 return return_shape; in TryGetSubshape()
936 Shape* return_shape = shape; in GetMutableSubshape() local
[all …]
/external/tensorflow/tensorflow/compiler/xla/service/
Dservice.cc921 Shape return_shape; in TransferToClient() local
923 return_shape = Shape(arg->shape_with_layout()); in TransferToClient()
924 if (!LayoutUtil::HasLayout(return_shape)) { in TransferToClient()
928 return_shape = Shape(shaped_buffer->on_host_shape()); in TransferToClient()
939 if (LayoutUtil::LayoutsInShapesEqual(return_shape, result_literal.shape())) { in TransferToClient()
943 result_literal.Relayout(return_shape).ToProto(); in TransferToClient()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dir_emitter.cc121 const Shape& return_shape = computation->root_instruction()->shape(); in EmitThreadLocalFunctionEpilogue() local
123 if (ShapeUtil::IsScalar(return_shape)) { in EmitThreadLocalFunctionEpilogue()
129 CHECK(return_shape.IsTuple()); in EmitThreadLocalFunctionEpilogue()
131 llvm::Type* tuple_type = llvm_ir::ShapeToIrType(return_shape, module_); in EmitThreadLocalFunctionEpilogue()
135 for (int i = 0; i < return_shape.tuple_shapes_size(); i++) { in EmitThreadLocalFunctionEpilogue()
136 const Shape& element_shape = return_shape.tuple_shapes(i); in EmitThreadLocalFunctionEpilogue()
3236 const Shape& return_shape = callee.root_instruction()->shape(); in EmitThreadLocalCall() local
3237 bool is_scalar_return = ShapeUtil::IsScalar(return_shape); in EmitThreadLocalCall()
3239 return_shape.IsTuple() && in EmitThreadLocalCall()
3240 absl::c_all_of(return_shape.tuple_shapes(), [&](const Shape& shape) { in EmitThreadLocalCall()
[all …]
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/ir/
Dtf_ops_n_z.cc1320 SmallVector<int64_t, 4> return_shape(input_rank, ShapedType::kDynamicSize); in inferReturnTypes() local
1325 {RankedTensorType::get(return_shape, input_type.getElementType())}); in inferReturnTypes()
1331 return_shape[i] = input_shape[i]; in inferReturnTypes()
1347 return_shape[1 + i] = ShapedType::kDynamicSize; in inferReturnTypes()
1356 return_shape[1 + i] = (paddings_sum + input_shape[i + 1]) / block_shape_i; in inferReturnTypes()
1358 return_shape[0] = return_batch; in inferReturnTypes()
1362 {RankedTensorType::get(return_shape, input_type.getElementType())}); in inferReturnTypes()
Dtf_ops_a_m.cc1633 SmallVector<int64_t, 4> return_shape(num_dims, ShapedType::kDynamicSize); in inferConvReturnTypes() local
1637 return_shape[GetTensorBatchDimIndex(num_dims, format)] = in inferConvReturnTypes()
1641 return_shape[GetTensorFeatureDimIndex(num_dims, format)] = in inferConvReturnTypes()
1669 return_shape[dim] = expected_output_size; in inferConvReturnTypes()
1674 {RankedTensorType::get(return_shape, input_ty.getElementType())}); in inferConvReturnTypes()
/external/tensorflow/tensorflow/compiler/xrt/tests/
Draw_api_test.cc1793 xla::Shape return_shape = xla::ShapeUtil::MakeTupleShape( in TEST() local
1820 *shapes->mutable_result() = return_shape.ToProto(); in TEST()
1892 xla::Shape return_shape = xla::ShapeUtil::MakeTupleShape( in TEST() local
1919 *shapes->mutable_result() = return_shape.ToProto(); in TEST()
/external/tensorflow/tensorflow/compiler/mlir/xla/
Dmlir_hlo_to_hlo.cc1319 xla::Shape return_shape = xla::TypeToShape(ret.get().getType()); in Lower() local
1322 builder, returns[index], return_shape, shape_representation_fn_, in Lower()