Home
last modified time | relevance | path

Searched refs:buffer_shape (Results 1 – 4 of 4) sorted by relevance

/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dscatter_nd_op.cc31 Status ValidateUpdateShape(const TensorShape& buffer_shape, in ValidateUpdateShape() argument
49 ", buffer_shape: ", buffer_shape.DebugString(), in ValidateUpdateShape()
54 if (buffer_shape.dims() < in ValidateUpdateShape()
59 batch_dim + buffer_shape.dims() - num_index_dims) { in ValidateUpdateShape()
69 buffer_shape.dim_size(d + num_index_dims)) { in ValidateUpdateShape()
86 TensorShape buffer_shape; in Compile() local
87 OP_REQUIRES_OK(context, context->ConstantInputAsShape(2, &buffer_shape)); in Compile()
90 context, TensorShapeUtils::IsVectorOrHigher(buffer_shape), in Compile()
92 "got shape: ", buffer_shape.DebugString())); in Compile()
96 buffer_shape.num_elements() > 0 || (indices_shape.num_elements() == 0 && in Compile()
[all …]
Dsegment_reduction_ops.cc66 TensorShape buffer_shape = data_shape; in Compile() local
67 buffer_shape.RemoveDimRange(0, indices_shape.dims()); in Compile()
68 buffer_shape.InsertDim(0, num_segments); in Compile()
70 buffer_shape.dim_sizes()); in Compile()
/external/tensorflow/tensorflow/compiler/tf2xla/lib/
Dscatter.cc41 TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> buffer_shape, in XlaScatter()
50 xla::AsInt64Slice(buffer_shape->dimensions()); in XlaScatter()
58 if (num_index_dims > xla::ShapeUtil::Rank(*buffer_shape)) { in XlaScatter()
63 xla::ShapeUtil::HumanString(*buffer_shape), ")"); in XlaScatter()
81 if (xla::ShapeUtil::GetDimension(*buffer_shape, i) == 0) { in XlaScatter()
84 xla::ShapeUtil::HumanString(*buffer_shape)); in XlaScatter()
/external/tensorflow/tensorflow/compiler/tf2xla/
Dxla_compiler.cc275 xla::Shape buffer_shape; in XLAShapeForArgument() local
277 TensorShapeToXLAShape(arg.type, shape, &buffer_shape)); in XLAShapeForArgument()
279 {buffer_shape, xla::ShapeUtil::MakeShape(xla::S32, {})}); in XLAShapeForArgument()