| /external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/ |
| D | lower_general_dot.cc | 52 llvm::ArrayRef<int64_t> arg_shape, in TransposeReshape() argument
58 left_size *= arg_shape[dim]; in TransposeReshape()
63 right_size *= arg_shape[dim]; in TransposeReshape()
83 transposed_shape.push_back(arg_shape[val]); in TransposeReshape()
|
| /external/tensorflow/tensorflow/compiler/jit/ |
| D | shape_inference.cc | 83 const InferredShape& arg_shape = it->second; in PropagateShapes() local
87 if (arg_shape.handle_type != DT_INVALID) { in PropagateShapes()
90 arg_shape.handle_shape, &handle)); in PropagateShapes()
95 {handle, arg_shape.handle_type}}); in PropagateShapes()
100 context->MakeShapeFromPartialTensorShape(arg_shape.shape, &handle)); in PropagateShapes()
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | compile_mlir_util.cc | 72 TensorShape arg_shape; in GetTensorShapeFromXlaArgument() local
74 XLAShapeToTensorShape(absl::get<xla::Shape>(arg.shape), &arg_shape)); in GetTensorShapeFromXlaArgument()
75 return arg_shape; in GetTensorShapeFromXlaArgument()
506 for (const auto& arg_shape : arg_shapes) in CompileSerializedMlirToXlaHlo() local
507 tensor_or_resource_shapes.push_back({arg_shape}); in CompileSerializedMlirToXlaHlo()
538 TF_ASSIGN_OR_RETURN(TensorShape arg_shape, in RewriteWithArgs()
540 auto resource_shape = arg_shape.dim_sizes(); in RewriteWithArgs()
595 TF_ASSIGN_OR_RETURN(TensorShape arg_shape, in CompileGraphSetup()
597 arg_shapes.push_back({arg_shape, in CompileGraphSetup()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | shape_inference.cc | 391 const Shape* arg_shape = nullptr; in InferConcatOpShape() local
395 if (!arg_shape) { in InferConcatOpShape()
396 arg_shape = shape; in InferConcatOpShape()
397 element_type = arg_shape->element_type(); in InferConcatOpShape()
400 if (arg_shape->rank() != shape->rank()) { in InferConcatOpShape()
404 arg_shape->rank(), ShapeUtil::HumanString(*arg_shape), shape->rank(), in InferConcatOpShape()
407 if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(*arg_shape, *shape)) { in InferConcatOpShape()
410 PrimitiveType_Name(arg_shape->element_type()), in InferConcatOpShape()
413 for (int64 dimension_number = 0; dimension_number < arg_shape->rank(); in InferConcatOpShape()
415 if (arg_shape->dimensions(dimension_number) != in InferConcatOpShape()
[all …]
|
| D | hlo_evaluator.cc | 247 const auto& arg_shape = arg_literals[i]->shape(); in Evaluate() local
249 arg_shape)) { in Evaluate()
254 ShapeUtil::HumanStringWithLayout(arg_shape)); in Evaluate()
2366 const Shape& arg_shape = input_args[0]->shape(); in GenerateReduceOutputElement() local
2367 absl::Span<const int64> arg_dimensions = AsInt64Slice(arg_shape.dimensions()); in GenerateReduceOutputElement()
2387 arg_shape, base, arg_dim_counts, arg_dim_steps, reduction_step)); in GenerateReduceOutputElement()
2395 arg_shape, base, arg_dim_counts, arg_dim_steps, in GenerateReduceOutputElement()
2435 const Shape& arg_shape = input_args[0]->shape(); in HandleReduce() local
2442 absl::Span<const int64> arg_dimensions = AsInt64Slice(arg_shape.dimensions()); in HandleReduce()
|
| D | shape_inference_test.cc | 1228 Shape arg_shape = ShapeUtil::MakeShape(F32, {5, 3}); in TEST_F() local
1230 {&arg_shape, &f32_}, in TEST_F()
1239 Shape arg_shape = ShapeUtil::MakeShape(F32, {5, 3}); in TEST_F() local
1241 ShapeInference::InferReduceShape({&arg_shape, &f32_}, in TEST_F()
1250 Shape arg_shape = ShapeUtil::MakeShape(F32, {5, 3}); in TEST_F() local
1252 ShapeInference::InferReduceShape({&arg_shape, &f32_}, in TEST_F()
1261 Shape arg_shape = ShapeUtil::MakeShape(F32, {5, 3}); in TEST_F() local
1263 {&arg_shape, &f32_}, in TEST_F()
1628 Shape arg_shape = ShapeUtil::MakeShape(F32, {5, 3}); in TEST_F() local
1631 arg_shape, val_shape, /*dimension=*/0); in TEST_F()
[all …]
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | fake_quantize_ops.cc | 101 TF_ASSIGN_OR_RETURN(xla::Shape arg_shape, b->GetShape(arg)); in BuildFakeQuantCustomCall()
102 *arg_shape.mutable_layout() = in BuildFakeQuantCustomCall()
103 xla::LayoutUtil::MakeDescendingLayout(arg_shape.rank()); in BuildFakeQuantCustomCall()
104 arg_shapes.push_back(std::move(arg_shape)); in BuildFakeQuantCustomCall()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | client.cc | 235 for (const auto& arg_shape : argument_shapes) { in Compile() local
236 *request.add_input_shape_with_layout() = arg_shape.ToProto(); in Compile()
|
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | distributed_tpu_rewrite_pass.cc | 1768 InferredShape& arg_shape = in GetArgAndRetvalShapes() local
1772 arg_shape.shape = TensorShape(); // Variables are always scalars. in GetArgAndRetvalShapes()
1773 arg_shape.handle_shape = info->handle_shape; in GetArgAndRetvalShapes()
1774 arg_shape.handle_type = info->handle_type; in GetArgAndRetvalShapes()
1775 TF_RET_CHECK(arg_shape.handle_type != DT_INVALID) in GetArgAndRetvalShapes()
2364 const InferredShape& arg_shape = arg_shapes[i]; in BuildCompileNode() local
2367 TF_RET_CHECK(arg_shape.handle_type != DT_INVALID) << i; in BuildCompileNode()
2368 arg->set_dtype(arg_shape.handle_type); in BuildCompileNode()
2369 arg_shape.handle_shape.AsProto(arg->mutable_shape()); in BuildCompileNode()
2374 arg_shape.shape.AsProto(arg->mutable_shape()); in BuildCompileNode()
[all …]
|
| /external/tensorflow/tensorflow/compiler/mlir/tools/kernel_gen/tests/ |
| D | buffer_reuse.mlir | 353 %arg_shape : memref<?xindex>,
373 %result = memref_reshape %flat_result(%arg_shape)
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/ |
| D | mlir_hlo_to_hlo.cc | 1450 xla::Shape& arg_shape = arg_shapes->back(); in SetEntryTupleShapesAndLeafReplication() local
1453 tensorflow::XLAShapeToTensorShape(arg_shape, &arg_tensor_shape); in SetEntryTupleShapesAndLeafReplication()
1468 arg_shape = std::move(arg_shape_status.ValueOrDie()); in SetEntryTupleShapesAndLeafReplication()
1471 for (int i = 0, e = xla::ShapeUtil::GetLeafCount(arg_shape); i < e; ++i) in SetEntryTupleShapesAndLeafReplication()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/ |
| D | xla_compiler_test.cc | 1621 xla::Shape arg_shape = xla::ShapeUtil::MakeTupleShape(shapes); in TEST_F() local
1622 args[0].shape = arg_shape; in TEST_F()
|