Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Improve shape function for CudnnRNNParamsSize #21114

Merged
merged 5 commits into from
Sep 20, 2018
Merged

Improve shape function for CudnnRNNParamsSize #21114

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
27de8e7
Improve shape function for CudnnRNNParamsSize
yongtang Jul 25, 2018
01387cc
Add test cases for shape function of CudnnRNNParamsSize
yongtang Jul 25, 2018
d27b5a3
Pylint fix
yongtang Jul 25, 2018
c863279
Add additional unit test in c++ for cudnn_rnn_ops
yongtang Jul 25, 2018
b127c20
Fix clang-format issue in `Experimental clang-format Check`
yongtang Aug 10, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
25 changes: 25 additions & 0 deletions tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_ops_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -413,6 +413,31 @@ def testLSTMParamsSize(self):
self._testOneLSTMParamsSize(num_layers, num_units, input_size,
direction)

@unittest.skipUnless(test.is_built_with_cuda(),
"Test only applicable when running on GPUs")
def testLSTMParamsSizeShape(self):
with self.assertRaisesRegexp(
ValueError, "Shape must be rank 0 but is rank 1"):
model = _CreateModel(
cudnn_rnn_ops.CUDNN_LSTM,
constant_op.constant([4]), 200, 200,
direction=cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION)
params_size = model.params_size()
with self.assertRaisesRegexp(
ValueError, "Shape must be rank 0 but is rank 1"):
model = _CreateModel(
cudnn_rnn_ops.CUDNN_LSTM,
4, constant_op.constant([200]), 200,
direction=cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION)
params_size = model.params_size()
with self.assertRaisesRegexp(
ValueError, "Shape must be rank 0 but is rank 1"):
model = _CreateModel(
cudnn_rnn_ops.CUDNN_LSTM,
4, 200, constant_op.constant([200]),
direction=cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION)
params_size = model.params_size()


class CudnnRNNTestInference(TensorFlowTestCase):

Expand Down
9 changes: 6 additions & 3 deletions tensorflow/core/ops/cudnn_rnn_ops.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,6 @@ using shape_inference::DimensionHandle;
using shape_inference::InferenceContext;
using shape_inference::ShapeHandle;


REGISTER_OP("CudnnRNNParamsSize")
.Input("num_layers: int32")
.Input("num_units: int32")
Expand All @@ -52,11 +51,16 @@ REGISTER_OP("CudnnRNNParamsSize")
.Attr("seed2: int = 0")
.Output("params_size: S")
.SetShapeFn([](InferenceContext* c) {
ShapeHandle unused;
// num_layers, num_units, and input_size should be scalars.
TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &unused));
TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));

c->set_output(0, c->Vector(1));
return Status::OK();
});


REGISTER_OP("CudnnRNN")
.Input("input: T")
.Input("input_h: T")
Expand Down Expand Up @@ -248,7 +252,6 @@ REGISTER_OP("CudnnRNNParamsToCanonical")
return Status::OK();
});


REGISTER_OP("CudnnRNNCanonicalToParams")
.Input("num_layers: int32")
.Input("num_units: int32")
Expand Down
11 changes: 10 additions & 1 deletion tensorflow/core/ops/cudnn_rnn_ops_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,16 @@ namespace tensorflow {

TEST(CudnnRNNOpsTest, ParamsSize_ShapeFn) {
ShapeInferenceTestOp op("CudnnRNNParamsSize");
INFER_OK(op, "[1];[1];[1]", "[1]");
INFER_OK(op, "[];[];[]", "[1]");
INFER_OK(op, "?;[];[]", "[1]");
INFER_OK(op, "[];?;[]", "[1]");
INFER_OK(op, "[];[];?", "[1]");
INFER_OK(op, "[];?;?", "[1]");
INFER_OK(op, "?;?;?", "[1]");

INFER_ERROR("Shape must be rank 0 ", op, "[1,2];?;[]");
INFER_ERROR("Shape must be rank 0 ", op, "?;[2];[]");
INFER_ERROR("Shape must be rank 0 ", op, "?;?;[1]");
}

TEST(CudnnRNNOpsTest, ForwardLstm_ShapeFn) {
Expand Down