support 1D inputs in leaky relu

apache · Jul 21, 2018 · fcf0ae6 · fcf0ae6
1 parent 11b69f2
commit fcf0ae6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/src/operator/leaky_relu-inl.h b/src/operator/leaky_relu-inl.h
@@ -94,7 +94,7 @@ class LeakyReLUOp : public Operator {
  Tensor<xpu, 3, DType> out;
  Tensor<xpu, 3, DType> mask;
  int n = in_data[leakyrelu::kData].shape_[0];
- int k = in_data[leakyrelu::kData].shape_[1];
+ int k = (in_data[leakyrelu::kData].ndim() > 1) ? in_data[leakyrelu::kData].shape_[1] : 1;
  Shape<3> dshape = Shape3(n, k, in_data[leakyrelu::kData].Size()/n/k);
  data = in_data[leakyrelu::kData].get_with_shape<xpu, 3, DType>(dshape, s);
  out = out_data[leakyrelu::kOut].get_with_shape<xpu, 3, DType>(dshape, s);
@@ -207,7 +207,7 @@ class LeakyReLUOp : public Operator {
  Tensor<xpu, 3, DType> grad;
  Tensor<xpu, 3, DType> mask;
  int n = out_grad[leakyrelu::kOut].shape_[0];
- int k = out_grad[leakyrelu::kOut].shape_[1];
+ int k = (out_grad[leakyrelu::kOut].ndim() > 1) ? out_grad[leakyrelu::kOut].shape_[1] : 1;
  Shape<3> dshape = Shape3(n, k, out_grad[leakyrelu::kOut].Size()/n/k);
  grad = out_grad[leakyrelu::kOut].get_with_shape<xpu, 3, DType>(dshape, s);
  gdata = in_grad[leakyrelu::kData].get_with_shape<xpu, 3, DType>(dshape, s);

diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py
@@ -748,7 +748,8 @@ def fleaky_relu_grad(grad, x, y, act_type, slope=0.25):
  elif act_type == 'leaky':
  out[neg_indices] = slope
  return out * grad
- shape = (3, 4)
+ ndim = random.randint(1, 4)
+ shape = rand_shape_nd(ndim)
  x = mx.symbol.Variable("x")
  slp = 0.25
  for dtype in [np.float16, np.float32, np.float64]: