Add slice_axis methods to Distribution (awslabs#397)

* add slice, slice_axis to Distribution, add tests

* remove slice method

* add header

* address comment

src/gluonts/distribution/binned.py

@@ -174,6 +174,10 @@ def s(bin_probs):
 
  return _sample_multiple(s, self.bin_probs, num_samples=num_samples)
 
+ @property
+ def args(self) -> List:
+ return [self.bin_probs, self.bin_centers]
+
 
 class BinnedArgs(gluon.HybridBlock):
  def __init__(

src/gluonts/distribution/distribution.py

@@ -253,6 +253,19 @@ def quantile(self, level: Tensor) -> Tensor:
  """
  raise NotImplementedError()
 
+ def slice_axis(
+ self, axis: int, begin: int, end: Optional[int]
+ ) -> "Distribution":
+ """
+ Construct a new distribution by slicing all constructor arguments
+ as specified by the provided bounds. Relies on ``mx.nd.slice_axis``.
+ """
+ sliced_distr = self.__class__(
+ *[arg.slice_axis(axis, begin, end) for arg in self.args]
+ )
+ assert isinstance(sliced_distr, type(self))
+ return sliced_distr
+
 
 def _expand_param(p: Tensor, num_samples: Optional[int] = None) -> Tensor:
  """

src/gluonts/distribution/gaussian.py

@@ -13,7 +13,7 @@
 
 # Standard library imports
 import math
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional, Tuple, List
 
 # First-party imports
 from gluonts.model.common import Tensor
@@ -116,6 +116,10 @@ def quantile(self, level: Tensor) -> Tensor:
  ),
  )
 
+ @property
+ def args(self) -> List:
+ return [self.mu, self.sigma]
+
 
 class GaussianOutput(DistributionOutput):
  args_dim: Dict[str, int] = {"mu": 1, "sigma": 1}

src/gluonts/distribution/laplace.py

@@ -12,7 +12,7 @@
 # permissions and limitations under the License.
 
 # Standard library imports
-from typing import Dict, Tuple
+from typing import Dict, Tuple, List
 
 # First-party imports
 from gluonts.model.common import Tensor
@@ -100,6 +100,10 @@ def quantile(self, level: Tensor) -> Tensor:
 
  return F.broadcast_add(self.mu, F.broadcast_mul(self.b, u))
 
+ @property
+ def args(self) -> List:
+ return [self.mu, self.b]
+
 
 class LaplaceOutput(DistributionOutput):
  args_dim: Dict[str, int] = {"mu": 1, "b": 1}

src/gluonts/distribution/neg_binomial.py

@@ -12,7 +12,7 @@
 # permissions and limitations under the License.
 
 # Standard library imports
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional, Tuple, List
 
 # First-party imports
 from gluonts.model.common import Tensor
@@ -93,6 +93,10 @@ def s(mu: Tensor, alpha: Tensor) -> Tensor:
  s, mu=self.mu, alpha=self.alpha, num_samples=num_samples
  )
 
+ @property
+ def args(self) -> List:
+ return [self.mu, self.alpha]
+
 
 class NegativeBinomialOutput(DistributionOutput):
  args_dim: Dict[str, int] = {"mu": 1, "alpha": 1}

src/gluonts/distribution/student_t.py

@@ -13,7 +13,7 @@
 
 # Standard library imports
 import math
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional, Tuple, List
 
 # First-party imports
 from gluonts.model.common import Tensor
@@ -111,6 +111,10 @@ def s(mu: Tensor, sigma: Tensor, nu: Tensor) -> Tensor:
  num_samples=num_samples,
  )
 
+ @property
+ def args(self) -> List:
+ return [self.mu, self.sigma, self.nu]
+
 
 class StudentTOutput(DistributionOutput):
  args_dim: Dict[str, int] = {"mu": 1, "sigma": 1, "nu": 1}

src/gluonts/distribution/uniform.py

@@ -12,7 +12,7 @@
 # permissions and limitations under the License.
 
 # Standard library imports
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional, Tuple, List
 
 # First-party imports
 from gluonts.model.common import Tensor
@@ -100,6 +100,10 @@ def quantile(self, level: Tensor) -> Tensor:
  F.broadcast_mul(self.high - self.low, level), self.low
  )
 
+ @property
+ def args(self) -> List:
+ return [self.low, self.high]
+
 
 class UniformOutput(DistributionOutput):
  args_dim: Dict[str, int] = {"low": 1, "width": 1}

test/distribution/test_distribution_slice.py

@@ -0,0 +1,40 @@
+# Copyright 2018 Amazon.com, Inc. or its affiliates. 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.
+# A copy of the License is located at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the "license" file accompanying this file. This file 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.
+
+import pytest
+
+import mxnet as mx
+
+from gluonts.distribution.gaussian import Gaussian
+
+DISTR_SHAPE = (3, 4)
+
+DISTR_CASES = [
+ Gaussian(
+ mu=mx.nd.random.normal(shape=DISTR_SHAPE),
+ sigma=mx.nd.random.uniform(shape=DISTR_SHAPE),
+ )
+]
+
+SLICE_AXIS_CASES = [[(0, 0, None), 3], [(0, 1, 3), 2], [(1, -1, None), 1]]
+
+
+@pytest.mark.parametrize(
+ "slice_axis_args, expected_axis_length", SLICE_AXIS_CASES
+)
+@pytest.mark.parametrize("distr", DISTR_CASES)
+def test_distr_slice_axis(distr, slice_axis_args, expected_axis_length):
+ axis, begin, end = slice_axis_args
+ distr_sliced = distr.slice_axis(axis, begin, end)
+
+ assert distr_sliced.batch_shape[axis] == expected_axis_length