experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add XLA kernel for tf.ensure_shape

Browse Source 
This PR is part of the effort to fix 34363 and a prerequisite for PR 34399 which will resolve 34363.

In order for the PR 34399 to pass all tests, an XLA kernel for tf.ensure_shape need
to be added.

Once this PR is merged, PR 34399 will be re-opened to address 34363.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
This commit is contained in:
Yong Tang 2020-04-14 16:15:29 +00:00
parent 20a26f65d0
commit da3c1f4b01
2 changed files with 55 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorflow/compiler/tf2xla/kernels/BUILD
View File

@ -38,6 +38,7 @@ tf_kernel_library(
"einsum_op.cc",
"elu_op.cc",
"elu_op.h",
"ensure_shape_op.cc",
"empty_op.cc",
"extract_image_patches_op.cc",
"fake_param_op.cc",

54
tensorflow/compiler/tf2xla/kernels/ensure_shape_op.cc Normal file
View File

@ -0,0 +1,54 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// XLA-specific ensure_shape Op.
#include "tensorflow/compiler/tf2xla/type_util.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/compiler/xla/literal.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
namespace tensorflow {
namespace {
class EnsureShapeOp : public XlaOpKernel {
public:
explicit EnsureShapeOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
OP_REQUIRES_OK(ctx, ctx->GetAttr("shape", &expected_shape_));
}
void Compile(XlaOpKernelContext* ctx) override {
const TensorShape shape = ctx->InputShape(0);
// valiate shape
OP_REQUIRES(ctx, expected_shape_.IsCompatibleWith(shape),
errors::InvalidArgument("Shape of tensor ", this->def().input(0), " ", shape.DebugString(), " is not compatible with expected shape ", expected_shape_.DebugString(), "."));
// If shape matches, outputs the tensor.
ctx->SetOutput(0, ctx->Input(0));
}
private:
PartialTensorShape expected_shape_;
};
REGISTER_XLA_OP(Name("EnsureShape"), EnsureShapeOp);
} // namespace
} // namespace tensorflow