Searched refs:arg_shape (Results 1 – 7 of 7) sorted by relevance
| /external/tensorflow/tensorflow/compiler/jit/ |
| D | shape_inference.cc | 72 const InferredShape& arg_shape = it->second; in PropagateShapes() local
76 if (arg_shape.handle_type != DT_INVALID) { in PropagateShapes()
79 arg_shape.handle_shape, &handle)); in PropagateShapes()
84 {handle, arg_shape.handle_type}}); in PropagateShapes()
89 context->MakeShapeFromPartialTensorShape(arg_shape.shape, &handle)); in PropagateShapes()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | index_ops_cpu.cc | 112 xla::Shape arg_shape = shape_status.ConsumeValueOrDie(); in Compile() local
113 *arg_shape.mutable_layout() = in Compile()
114 xla::LayoutUtil::MakeDescendingLayout(arg_shape.rank()); in Compile()
115 arg_shapes.push_back(std::move(arg_shape)); in Compile()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | shape_inference.cc | 357 const Shape* arg_shape = nullptr; in InferConcatOpShape() local
361 if (!arg_shape) { in InferConcatOpShape()
362 arg_shape = shape; in InferConcatOpShape()
363 element_type = arg_shape->element_type(); in InferConcatOpShape()
366 if (arg_shape->rank() != shape->rank()) { in InferConcatOpShape()
370 arg_shape->rank(), ShapeUtil::HumanString(*arg_shape), shape->rank(), in InferConcatOpShape()
373 if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(*arg_shape, *shape)) { in InferConcatOpShape()
376 PrimitiveType_Name(arg_shape->element_type()), in InferConcatOpShape()
379 for (int64 dimension_number = 0; dimension_number < arg_shape->rank(); in InferConcatOpShape()
381 if (arg_shape->dimensions(dimension_number) != in InferConcatOpShape()
[all …]
|
| D | shape_inference_test.cc | 802 Shape arg_shape = ShapeUtil::MakeShape(F32, {5, 3}); in TEST_F() local
804 {&arg_shape, &f32_}, in TEST_F()
813 Shape arg_shape = ShapeUtil::MakeShape(F32, {5, 3}); in TEST_F() local
815 ShapeInference::InferReduceShape({&arg_shape, &f32_}, in TEST_F()
824 Shape arg_shape = ShapeUtil::MakeShape(F32, {5, 3}); in TEST_F() local
826 ShapeInference::InferReduceShape({&arg_shape, &f32_}, in TEST_F()
|
| D | hlo_evaluator.cc | 245 const auto& arg_shape = arg_literals[i]->shape(); in Evaluate() local
246 if (!ShapeUtil::Equal(computation_shape, arg_shape)) { in Evaluate()
251 ShapeUtil::HumanString(arg_shape)); in Evaluate()
|
| D | hlo_evaluator_typed_visitor.h | 1717 const Shape& arg_shape = arg_literals[0]->shape(); in HandleReduce() local
1721 const auto arg_dimensions = AsInt64Slice(arg_shape.dimensions()); in HandleReduce()
1826 arg_shape, base, arg_dim_counts, arg_dim_steps, func); in HandleReduce()
|
| /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()
|