We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Binned distribution gives wrong quantiles if the probability of first bin is nonzero
## Description
While expanding bin centers, [the code](https://github.com/awslabs/gluon-ts/blob/a9cde7bf92178b4cf8010a4c8600940b595f2ae3/src/gluonts/distribution/binned.py#L154) incorrectly adds the probabilities of the first bin to the bin centers ([`zeros_cdf` has the probabilities of the initial bin for all batch elements/time points](https://github.com/awslabs/gluon-ts/blob/a9cde7bf92178b4cf8010a4c8600940b595f2ae3/src/gluonts/distribution/binned.py#L138)).

Also the [index](https://github.com/awslabs/gluon-ts/blob/a9cde7bf92178b4cf8010a4c8600940b595f2ae3/src/gluonts/distribution/binned.py#L146) returned is always zero and hence the quantile returned is the first bin + probability of the first bin, unless the initial bin probability is zero which is always the case in [tests](https://github.com/awslabs/gluon-ts/blob/a9cde7bf92178b4cf8010a4c8600940b595f2ae3/test/distribution/test_distribution_sampling.py#L84).

## To Reproduce
```
In [1]: import mxnet as mx                                                                                                                                                                     
from
In [2]: from gluonts.distribution import Binned                                                                                                                                                

In [3]: binned = Binned(bin_probs=mx.nd.array([[0.5, 0.2, 0.15, 0.15]]), bin_centers=mx.nd.array([[1e-4, 1e-3, 1e-2, 1e-1]]))                                                                  

In [4]: binned.quantile(mx.nd.array([0.1, 0.5, 0.7, 0.9, 0.999]))                                                                                                                              
Out[4]: 

[[0.5001]
 [0.5001]
 [0.5001]
 [0.5001]
 [0.5001]]
<NDArray 5x1 @cpu(0)>

```
## Fix
Replacing `broadcast_add` by `broadcast_mul` [here](https://github.com/awslabs/gluon-ts/blob/a9cde7bf92178b4cf8010a4c8600940b595f2ae3/src/gluonts/distribution/binned.py#L138) seems to solve both  problems. I think `zeros_cdf` was meant to be zeros not the probabilities of the first bin.

After fix:

```
In [4]: binned.quantile(mx.nd.array([0.1, 0.5, 0.7, 0.9, 0.999]))                                                                                                                              
Out[4]: 

[[1.e-04]
 [1.e-03]
 [1.e-02]
 [1.e-01]
 [1.e-01]]
<NDArray 5x1 @cpu(0)>

```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/gluonts/distribution/binned.py`
Content:
```
1 # Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License").
4 # You may not use this file except in compliance with the License.
5 # A copy of the License is located at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # or in the "license" file accompanying this file. This file is distributed
10 # on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
11 # express or implied. See the License for the specific language governing
12 # permissions and limitations under the License.
13 
14 # Standard library imports
15 from typing import Tuple, List
16 
17 # Third-party imports
18 import mxnet as mx
19 from mxnet import gluon
20 import numpy as np
21 
22 # First-party imports
23 from gluonts.core.component import validated
24 from gluonts.model.common import Tensor
25 
26 # Relative imports
27 from .distribution import Distribution, _sample_multiple, getF
28 from .distribution_output import DistributionOutput
29 
30 
31 class Binned(Distribution):
32     r"""
33     A binned distribution defined by a set of bins via
34     bin centers and bin probabilities.
35 
36     Parameters
37     ----------
38     bin_probs
39         Tensor containing the bin probabilities, of shape `(*batch_shape, num_bins)`.
40     bin_centers
41         Tensor containing the bin centers, of shape `(*batch_shape, num_bins)`.
42     F
43     """
44 
45     is_reparameterizable = False
46 
47     @validated()
48     def __init__(self, bin_probs: Tensor, bin_centers: Tensor, F=None) -> None:
49         self.bin_centers = bin_centers
50         self.bin_probs = bin_probs
51         self.F = F if F else getF(bin_probs)
52 
53         self.bin_edges = Binned._compute_edges(self.F, bin_centers)
54 
55     @staticmethod
56     def _compute_edges(F, bin_centers: Tensor) -> Tensor:
57         r"""
58         Computes the edges of the bins based on the centers. The first and last edge are set to :math:`10^{-10}` and
59         :math:`10^{10}`, repsectively.
60 
61         Parameters
62         ----------
63         F
64         bin_centers
65             Tensor of shape `(*batch_shape, num_bins)`.
66 
67         Returns
68         -------
69         Tensor
70             Tensor of shape (*batch.shape, num_bins+1)
71         """
72 
73         low = (
74             F.zeros_like(bin_centers.slice_axis(axis=-1, begin=0, end=1))
75             - 1.0e10
76         )
77         high = (
78             F.zeros_like(bin_centers.slice_axis(axis=-1, begin=0, end=1))
79             + 1.0e10
80         )
81 
82         means = (
83             bin_centers.slice_axis(axis=-1, begin=1, end=None)
84             + bin_centers.slice_axis(axis=-1, begin=0, end=-1)
85         ) / 2.0
86 
87         return F.concat(low, means, high, dim=-1)
88 
89     @property
90     def batch_shape(self) -> Tuple:
91         return self.bin_probs.shape[:-1]
92 
93     @property
94     def event_shape(self) -> Tuple:
95         return ()
96 
97     @property
98     def event_dim(self) -> int:
99         return 0
100 
101     @property
102     def mean(self):
103         return (self.bin_probs * self.bin_centers).sum(axis=-1)
104 
105     @property
106     def stddev(self):
107         Ex2 = (self.bin_probs * self.bin_centers.square()).sum(axis=-1)
108         return (Ex2 - self.mean.square()).sqrt()
109 
110     def log_prob(self, x):
111         F = self.F
112         x = x.expand_dims(axis=-1)
113         # TODO: when mxnet has searchsorted replace this
114         left_edges = self.bin_edges.slice_axis(axis=-1, begin=0, end=-1)
115         right_edges = self.bin_edges.slice_axis(axis=-1, begin=1, end=None)
116         mask = F.broadcast_lesser_equal(left_edges, x) * F.broadcast_lesser(
117             x, right_edges
118         )
119         return F.broadcast_mul(self.bin_probs.log(), mask).sum(axis=-1)
120 
121     def cdf(self, x: Tensor) -> Tensor:
122         F = self.F
123         x = x.expand_dims(axis=-1)
124         # left_edges = self.bin_edges.slice_axis(axis=-1, begin=0, end=-1)
125         mask = F.broadcast_lesser_equal(self.bin_centers, x)
126         return F.broadcast_mul(self.bin_probs, mask).sum(axis=-1)
127 
128     def quantile(self, level: Tensor) -> Tensor:
129         F = self.F
130 
131         probs = self.bin_probs.swapaxes(0, 1)  # (num_bins, batch)
132         zeros_batch_size = F.slice_axis(probs, axis=0, begin=0, end=1).squeeze(
133             axis=0
134         )  # (batch_size,)
135 
136         level = level.expand_dims(axis=0)
137         # cdf shape (batch_size, levels)
138         zeros_cdf = F.broadcast_add(
139             zeros_batch_size.expand_dims(axis=1), level.zeros_like()
140         )
141         start_state = (zeros_cdf, zeros_cdf.astype("int32"))
142 
143         def step(p, state):
144             cdf, idx = state
145             cdf = F.broadcast_add(cdf, p.expand_dims(axis=1))
146             idx = F.where(F.broadcast_greater(cdf, level), idx, idx + 1)
147             return zeros_batch_size, (cdf, idx)
148 
149         _, states = F.contrib.foreach(step, probs, start_state)
150         _, idx = states
151 
152         # expand centers to shape (batch, levels, num_bins)
153         # so we can use pick with idx.shape = (batch, levels)
154         centers_expanded = F.broadcast_add(
155             self.bin_centers.expand_dims(axis=1),
156             zeros_cdf.expand_dims(axis=-1),
157         )
158         a = centers_expanded.pick(idx, axis=-1)
159         return a.swapaxes(0, 1)
160 
161     def sample(self, num_samples=None, dtype=np.float32):
162         def s(bin_probs):
163             F = self.F
164             indices = F.sample_multinomial(bin_probs)
165             if num_samples is None:
166                 return self.bin_centers.pick(indices, -1).reshape_like(
167                     F.zeros_like(indices.astype("float32"))
168                 )
169             else:
170                 return F.repeat(
171                     F.expand_dims(self.bin_centers, axis=0),
172                     repeats=num_samples,
173                     axis=0,
174                 ).pick(indices, -1)
175 
176         return _sample_multiple(s, self.bin_probs, num_samples=num_samples)
177 
178     @property
179     def args(self) -> List:
180         return [self.bin_probs, self.bin_centers]
181 
182 
183 class BinnedArgs(gluon.HybridBlock):
184     def __init__(
185         self, num_bins: int, bin_centers: mx.nd.NDArray, **kwargs
186     ) -> None:
187         super().__init__(**kwargs)
188         self.num_bins = num_bins
189         with self.name_scope():
190             self.bin_centers = self.params.get_constant(
191                 "bin_centers", bin_centers
192             )
193 
194             # needs to be named self.proj for consistency with the
195             # ArgProj class and the inference tests
196             self.proj = gluon.nn.HybridSequential()
197             self.proj.add(
198                 gluon.nn.Dense(
199                     self.num_bins,
200                     prefix="binproj",
201                     flatten=False,
202                     weight_initializer=mx.init.Xavier(),
203                 )
204             )
205             self.proj.add(gluon.nn.HybridLambda("softmax"))
206 
207     def hybrid_forward(
208         self, F, x: Tensor, bin_centers: Tensor
209     ) -> Tuple[Tensor, Tensor]:
210         ps = self.proj(x)
211         reshaped_probs = ps.reshape(shape=(-2, -1, self.num_bins), reverse=1)
212         bin_centers = F.broadcast_add(bin_centers, ps.zeros_like())
213         return reshaped_probs, bin_centers
214 
215 
216 class BinnedOutput(DistributionOutput):
217     distr_cls: type = Binned
218 
219     @validated()
220     def __init__(self, bin_centers: mx.nd.NDArray) -> None:
221         self.bin_centers = bin_centers
222         self.num_bins = self.bin_centers.shape[0]
223         assert len(self.bin_centers.shape) == 1
224 
225     def get_args_proj(self, *args, **kwargs) -> gluon.nn.HybridBlock:
226         return BinnedArgs(self.num_bins, self.bin_centers)
227 
228     def distribution(self, args, scale=None) -> Binned:
229         probs = args[0]
230         bin_centers = args[1]
231         F = getF(probs)
232 
233         bin_centers = F.broadcast_mul(bin_centers, F.ones_like(probs))
234 
235         if scale is not None:
236             bin_centers = F.broadcast_mul(
237                 bin_centers, scale.expand_dims(axis=-1)
238             )
239 
240         return Binned(probs, bin_centers)
241 
242     @property
243     def event_shape(self) -> Tuple:
244         return ()
245
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/src/gluonts/distribution/binned.py b/src/gluonts/distribution/binned.py
--- a/src/gluonts/distribution/binned.py
+++ b/src/gluonts/distribution/binned.py
@@ -129,9 +129,10 @@
         F = self.F
 
         probs = self.bin_probs.swapaxes(0, 1)  # (num_bins, batch)
-        zeros_batch_size = F.slice_axis(probs, axis=0, begin=0, end=1).squeeze(
-            axis=0
-        )  # (batch_size,)
+        # (batch_size,)
+        zeros_batch_size = F.zeros_like(
+            F.slice_axis(probs, axis=0, begin=0, end=1).squeeze(axis=0)
+        )
 
         level = level.expand_dims(axis=0)
         # cdf shape (batch_size, levels)

{"golden_diff": "diff --git a/src/gluonts/distribution/binned.py b/src/gluonts/distribution/binned.py\n--- a/src/gluonts/distribution/binned.py\n+++ b/src/gluonts/distribution/binned.py\n@@ -129,9 +129,10 @@\n         F = self.F\n \n         probs = self.bin_probs.swapaxes(0, 1)  # (num_bins, batch)\n-        zeros_batch_size = F.slice_axis(probs, axis=0, begin=0, end=1).squeeze(\n-            axis=0\n-        )  # (batch_size,)\n+        # (batch_size,)\n+        zeros_batch_size = F.zeros_like(\n+            F.slice_axis(probs, axis=0, begin=0, end=1).squeeze(axis=0)\n+        )\n \n         level = level.expand_dims(axis=0)\n         # cdf shape (batch_size, levels)\n", "issue": "Binned distribution gives wrong quantiles if the probability of first bin is nonzero\n## Description\r\nWhile expanding bin centers, [the code](https://github.com/awslabs/gluon-ts/blob/a9cde7bf92178b4cf8010a4c8600940b595f2ae3/src/gluonts/distribution/binned.py#L154) incorrectly adds the probabilities of the first bin to the bin centers ([`zeros_cdf` has the probabilities of the initial bin for all batch elements/time points](https://github.com/awslabs/gluon-ts/blob/a9cde7bf92178b4cf8010a4c8600940b595f2ae3/src/gluonts/distribution/binned.py#L138)).\r\n\r\nAlso the [index](https://github.com/awslabs/gluon-ts/blob/a9cde7bf92178b4cf8010a4c8600940b595f2ae3/src/gluonts/distribution/binned.py#L146) returned is always zero and hence the quantile returned is the first bin + probability of the first bin, unless the initial bin probability is zero which is always the case in [tests](https://github.com/awslabs/gluon-ts/blob/a9cde7bf92178b4cf8010a4c8600940b595f2ae3/test/distribution/test_distribution_sampling.py#L84).\r\n\r\n## To Reproduce\r\n```\r\nIn [1]: import mxnet as mx                                                                                                                                                                     \r\nfrom\r\nIn [2]: from gluonts.distribution import Binned                                                                                                                                                \r\n\r\nIn [3]: binned = Binned(bin_probs=mx.nd.array([[0.5, 0.2, 0.15, 0.15]]), bin_centers=mx.nd.array([[1e-4, 1e-3, 1e-2, 1e-1]]))                                                                  \r\n\r\nIn [4]: binned.quantile(mx.nd.array([0.1, 0.5, 0.7, 0.9, 0.999]))                                                                                                                              \r\nOut[4]: \r\n\r\n[[0.5001]\r\n [0.5001]\r\n [0.5001]\r\n [0.5001]\r\n [0.5001]]\r\n<NDArray 5x1 @cpu(0)>\r\n\r\n```\r\n## Fix\r\nReplacing `broadcast_add` by `broadcast_mul` [here](https://github.com/awslabs/gluon-ts/blob/a9cde7bf92178b4cf8010a4c8600940b595f2ae3/src/gluonts/distribution/binned.py#L138) seems to solve both  problems. I think `zeros_cdf` was meant to be zeros not the probabilities of the first bin.\r\n\r\nAfter fix:\r\n\r\n```\r\nIn [4]: binned.quantile(mx.nd.array([0.1, 0.5, 0.7, 0.9, 0.999]))                                                                                                                              \r\nOut[4]: \r\n\r\n[[1.e-04]\r\n [1.e-03]\r\n [1.e-02]\r\n [1.e-01]\r\n [1.e-01]]\r\n<NDArray 5x1 @cpu(0)>\r\n\r\n```\n", "before_files": [{"content": "# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\").\n# You may not use this file except in compliance with the License.\n# A copy of the License is located at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# or in the \"license\" file accompanying this file. This file is distributed\n# on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either\n# express or implied. See the License for the specific language governing\n# permissions and limitations under the License.\n\n# Standard library imports\nfrom typing import Tuple, List\n\n# Third-party imports\nimport mxnet as mx\nfrom mxnet import gluon\nimport numpy as np\n\n# First-party imports\nfrom gluonts.core.component import validated\nfrom gluonts.model.common import Tensor\n\n# Relative imports\nfrom .distribution import Distribution, _sample_multiple, getF\nfrom .distribution_output import DistributionOutput\n\n\nclass Binned(Distribution):\n    r\"\"\"\n    A binned distribution defined by a set of bins via\n    bin centers and bin probabilities.\n\n    Parameters\n    ----------\n    bin_probs\n        Tensor containing the bin probabilities, of shape `(*batch_shape, num_bins)`.\n    bin_centers\n        Tensor containing the bin centers, of shape `(*batch_shape, num_bins)`.\n    F\n    \"\"\"\n\n    is_reparameterizable = False\n\n    @validated()\n    def __init__(self, bin_probs: Tensor, bin_centers: Tensor, F=None) -> None:\n        self.bin_centers = bin_centers\n        self.bin_probs = bin_probs\n        self.F = F if F else getF(bin_probs)\n\n        self.bin_edges = Binned._compute_edges(self.F, bin_centers)\n\n    @staticmethod\n    def _compute_edges(F, bin_centers: Tensor) -> Tensor:\n        r\"\"\"\n        Computes the edges of the bins based on the centers. The first and last edge are set to :math:`10^{-10}` and\n        :math:`10^{10}`, repsectively.\n\n        Parameters\n        ----------\n        F\n        bin_centers\n            Tensor of shape `(*batch_shape, num_bins)`.\n\n        Returns\n        -------\n        Tensor\n            Tensor of shape (*batch.shape, num_bins+1)\n        \"\"\"\n\n        low = (\n            F.zeros_like(bin_centers.slice_axis(axis=-1, begin=0, end=1))\n            - 1.0e10\n        )\n        high = (\n            F.zeros_like(bin_centers.slice_axis(axis=-1, begin=0, end=1))\n            + 1.0e10\n        )\n\n        means = (\n            bin_centers.slice_axis(axis=-1, begin=1, end=None)\n            + bin_centers.slice_axis(axis=-1, begin=0, end=-1)\n        ) / 2.0\n\n        return F.concat(low, means, high, dim=-1)\n\n    @property\n    def batch_shape(self) -> Tuple:\n        return self.bin_probs.shape[:-1]\n\n    @property\n    def event_shape(self) -> Tuple:\n        return ()\n\n    @property\n    def event_dim(self) -> int:\n        return 0\n\n    @property\n    def mean(self):\n        return (self.bin_probs * self.bin_centers).sum(axis=-1)\n\n    @property\n    def stddev(self):\n        Ex2 = (self.bin_probs * self.bin_centers.square()).sum(axis=-1)\n        return (Ex2 - self.mean.square()).sqrt()\n\n    def log_prob(self, x):\n        F = self.F\n        x = x.expand_dims(axis=-1)\n        # TODO: when mxnet has searchsorted replace this\n        left_edges = self.bin_edges.slice_axis(axis=-1, begin=0, end=-1)\n        right_edges = self.bin_edges.slice_axis(axis=-1, begin=1, end=None)\n        mask = F.broadcast_lesser_equal(left_edges, x) * F.broadcast_lesser(\n            x, right_edges\n        )\n        return F.broadcast_mul(self.bin_probs.log(), mask).sum(axis=-1)\n\n    def cdf(self, x: Tensor) -> Tensor:\n        F = self.F\n        x = x.expand_dims(axis=-1)\n        # left_edges = self.bin_edges.slice_axis(axis=-1, begin=0, end=-1)\n        mask = F.broadcast_lesser_equal(self.bin_centers, x)\n        return F.broadcast_mul(self.bin_probs, mask).sum(axis=-1)\n\n    def quantile(self, level: Tensor) -> Tensor:\n        F = self.F\n\n        probs = self.bin_probs.swapaxes(0, 1)  # (num_bins, batch)\n        zeros_batch_size = F.slice_axis(probs, axis=0, begin=0, end=1).squeeze(\n            axis=0\n        )  # (batch_size,)\n\n        level = level.expand_dims(axis=0)\n        # cdf shape (batch_size, levels)\n        zeros_cdf = F.broadcast_add(\n            zeros_batch_size.expand_dims(axis=1), level.zeros_like()\n        )\n        start_state = (zeros_cdf, zeros_cdf.astype(\"int32\"))\n\n        def step(p, state):\n            cdf, idx = state\n            cdf = F.broadcast_add(cdf, p.expand_dims(axis=1))\n            idx = F.where(F.broadcast_greater(cdf, level), idx, idx + 1)\n            return zeros_batch_size, (cdf, idx)\n\n        _, states = F.contrib.foreach(step, probs, start_state)\n        _, idx = states\n\n        # expand centers to shape (batch, levels, num_bins)\n        # so we can use pick with idx.shape = (batch, levels)\n        centers_expanded = F.broadcast_add(\n            self.bin_centers.expand_dims(axis=1),\n            zeros_cdf.expand_dims(axis=-1),\n        )\n        a = centers_expanded.pick(idx, axis=-1)\n        return a.swapaxes(0, 1)\n\n    def sample(self, num_samples=None, dtype=np.float32):\n        def s(bin_probs):\n            F = self.F\n            indices = F.sample_multinomial(bin_probs)\n            if num_samples is None:\n                return self.bin_centers.pick(indices, -1).reshape_like(\n                    F.zeros_like(indices.astype(\"float32\"))\n                )\n            else:\n                return F.repeat(\n                    F.expand_dims(self.bin_centers, axis=0),\n                    repeats=num_samples,\n                    axis=0,\n                ).pick(indices, -1)\n\n        return _sample_multiple(s, self.bin_probs, num_samples=num_samples)\n\n    @property\n    def args(self) -> List:\n        return [self.bin_probs, self.bin_centers]\n\n\nclass BinnedArgs(gluon.HybridBlock):\n    def __init__(\n        self, num_bins: int, bin_centers: mx.nd.NDArray, **kwargs\n    ) -> None:\n        super().__init__(**kwargs)\n        self.num_bins = num_bins\n        with self.name_scope():\n            self.bin_centers = self.params.get_constant(\n                \"bin_centers\", bin_centers\n            )\n\n            # needs to be named self.proj for consistency with the\n            # ArgProj class and the inference tests\n            self.proj = gluon.nn.HybridSequential()\n            self.proj.add(\n                gluon.nn.Dense(\n                    self.num_bins,\n                    prefix=\"binproj\",\n                    flatten=False,\n                    weight_initializer=mx.init.Xavier(),\n                )\n            )\n            self.proj.add(gluon.nn.HybridLambda(\"softmax\"))\n\n    def hybrid_forward(\n        self, F, x: Tensor, bin_centers: Tensor\n    ) -> Tuple[Tensor, Tensor]:\n        ps = self.proj(x)\n        reshaped_probs = ps.reshape(shape=(-2, -1, self.num_bins), reverse=1)\n        bin_centers = F.broadcast_add(bin_centers, ps.zeros_like())\n        return reshaped_probs, bin_centers\n\n\nclass BinnedOutput(DistributionOutput):\n    distr_cls: type = Binned\n\n    @validated()\n    def __init__(self, bin_centers: mx.nd.NDArray) -> None:\n        self.bin_centers = bin_centers\n        self.num_bins = self.bin_centers.shape[0]\n        assert len(self.bin_centers.shape) == 1\n\n    def get_args_proj(self, *args, **kwargs) -> gluon.nn.HybridBlock:\n        return BinnedArgs(self.num_bins, self.bin_centers)\n\n    def distribution(self, args, scale=None) -> Binned:\n        probs = args[0]\n        bin_centers = args[1]\n        F = getF(probs)\n\n        bin_centers = F.broadcast_mul(bin_centers, F.ones_like(probs))\n\n        if scale is not None:\n            bin_centers = F.broadcast_mul(\n                bin_centers, scale.expand_dims(axis=-1)\n            )\n\n        return Binned(probs, bin_centers)\n\n    @property\n    def event_shape(self) -> Tuple:\n        return ()\n", "path": "src/gluonts/distribution/binned.py"}], "after_files": [{"content": "# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\").\n# You may not use this file except in compliance with the License.\n# A copy of the License is located at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# or in the \"license\" file accompanying this file. This file is distributed\n# on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either\n# express or implied. See the License for the specific language governing\n# permissions and limitations under the License.\n\n# Standard library imports\nfrom typing import Tuple, List\n\n# Third-party imports\nimport mxnet as mx\nfrom mxnet import gluon\nimport numpy as np\n\n# First-party imports\nfrom gluonts.core.component import validated\nfrom gluonts.model.common import Tensor\n\n# Relative imports\nfrom .distribution import Distribution, _sample_multiple, getF\nfrom .distribution_output import DistributionOutput\n\n\nclass Binned(Distribution):\n    r\"\"\"\n    A binned distribution defined by a set of bins via\n    bin centers and bin probabilities.\n\n    Parameters\n    ----------\n    bin_probs\n        Tensor containing the bin probabilities, of shape `(*batch_shape, num_bins)`.\n    bin_centers\n        Tensor containing the bin centers, of shape `(*batch_shape, num_bins)`.\n    F\n    \"\"\"\n\n    is_reparameterizable = False\n\n    @validated()\n    def __init__(self, bin_probs: Tensor, bin_centers: Tensor, F=None) -> None:\n        self.bin_centers = bin_centers\n        self.bin_probs = bin_probs\n        self.F = F if F else getF(bin_probs)\n\n        self.bin_edges = Binned._compute_edges(self.F, bin_centers)\n\n    @staticmethod\n    def _compute_edges(F, bin_centers: Tensor) -> Tensor:\n        r\"\"\"\n        Computes the edges of the bins based on the centers. The first and last edge are set to :math:`10^{-10}` and\n        :math:`10^{10}`, repsectively.\n\n        Parameters\n        ----------\n        F\n        bin_centers\n            Tensor of shape `(*batch_shape, num_bins)`.\n\n        Returns\n        -------\n        Tensor\n            Tensor of shape (*batch.shape, num_bins+1)\n        \"\"\"\n\n        low = (\n            F.zeros_like(bin_centers.slice_axis(axis=-1, begin=0, end=1))\n            - 1.0e10\n        )\n        high = (\n            F.zeros_like(bin_centers.slice_axis(axis=-1, begin=0, end=1))\n            + 1.0e10\n        )\n\n        means = (\n            bin_centers.slice_axis(axis=-1, begin=1, end=None)\n            + bin_centers.slice_axis(axis=-1, begin=0, end=-1)\n        ) / 2.0\n\n        return F.concat(low, means, high, dim=-1)\n\n    @property\n    def batch_shape(self) -> Tuple:\n        return self.bin_probs.shape[:-1]\n\n    @property\n    def event_shape(self) -> Tuple:\n        return ()\n\n    @property\n    def event_dim(self) -> int:\n        return 0\n\n    @property\n    def mean(self):\n        return (self.bin_probs * self.bin_centers).sum(axis=-1)\n\n    @property\n    def stddev(self):\n        Ex2 = (self.bin_probs * self.bin_centers.square()).sum(axis=-1)\n        return (Ex2 - self.mean.square()).sqrt()\n\n    def log_prob(self, x):\n        F = self.F\n        x = x.expand_dims(axis=-1)\n        # TODO: when mxnet has searchsorted replace this\n        left_edges = self.bin_edges.slice_axis(axis=-1, begin=0, end=-1)\n        right_edges = self.bin_edges.slice_axis(axis=-1, begin=1, end=None)\n        mask = F.broadcast_lesser_equal(left_edges, x) * F.broadcast_lesser(\n            x, right_edges\n        )\n        return F.broadcast_mul(self.bin_probs.log(), mask).sum(axis=-1)\n\n    def cdf(self, x: Tensor) -> Tensor:\n        F = self.F\n        x = x.expand_dims(axis=-1)\n        # left_edges = self.bin_edges.slice_axis(axis=-1, begin=0, end=-1)\n        mask = F.broadcast_lesser_equal(self.bin_centers, x)\n        return F.broadcast_mul(self.bin_probs, mask).sum(axis=-1)\n\n    def quantile(self, level: Tensor) -> Tensor:\n        F = self.F\n\n        probs = self.bin_probs.swapaxes(0, 1)  # (num_bins, batch)\n        # (batch_size,)\n        zeros_batch_size = F.zeros_like(\n            F.slice_axis(probs, axis=0, begin=0, end=1).squeeze(axis=0)\n        )\n\n        level = level.expand_dims(axis=0)\n        # cdf shape (batch_size, levels)\n        zeros_cdf = F.broadcast_add(\n            zeros_batch_size.expand_dims(axis=1), level.zeros_like()\n        )\n        start_state = (zeros_cdf, zeros_cdf.astype(\"int32\"))\n\n        def step(p, state):\n            cdf, idx = state\n            cdf = F.broadcast_add(cdf, p.expand_dims(axis=1))\n            idx = F.where(F.broadcast_greater(cdf, level), idx, idx + 1)\n            return zeros_batch_size, (cdf, idx)\n\n        _, states = F.contrib.foreach(step, probs, start_state)\n        _, idx = states\n\n        # expand centers to shape (batch, levels, num_bins)\n        # so we can use pick with idx.shape = (batch, levels)\n        centers_expanded = F.broadcast_add(\n            self.bin_centers.expand_dims(axis=1),\n            zeros_cdf.expand_dims(axis=-1),\n        )\n        a = centers_expanded.pick(idx, axis=-1)\n        return a.swapaxes(0, 1)\n\n    def sample(self, num_samples=None, dtype=np.float32):\n        def s(bin_probs):\n            F = self.F\n            indices = F.sample_multinomial(bin_probs)\n            if num_samples is None:\n                return self.bin_centers.pick(indices, -1).reshape_like(\n                    F.zeros_like(indices.astype(\"float32\"))\n                )\n            else:\n                return F.repeat(\n                    F.expand_dims(self.bin_centers, axis=0),\n                    repeats=num_samples,\n                    axis=0,\n                ).pick(indices, -1)\n\n        return _sample_multiple(s, self.bin_probs, num_samples=num_samples)\n\n    @property\n    def args(self) -> List:\n        return [self.bin_probs, self.bin_centers]\n\n\nclass BinnedArgs(gluon.HybridBlock):\n    def __init__(\n        self, num_bins: int, bin_centers: mx.nd.NDArray, **kwargs\n    ) -> None:\n        super().__init__(**kwargs)\n        self.num_bins = num_bins\n        with self.name_scope():\n            self.bin_centers = self.params.get_constant(\n                \"bin_centers\", bin_centers\n            )\n\n            # needs to be named self.proj for consistency with the\n            # ArgProj class and the inference tests\n            self.proj = gluon.nn.HybridSequential()\n            self.proj.add(\n                gluon.nn.Dense(\n                    self.num_bins,\n                    prefix=\"binproj\",\n                    flatten=False,\n                    weight_initializer=mx.init.Xavier(),\n                )\n            )\n            self.proj.add(gluon.nn.HybridLambda(\"softmax\"))\n\n    def hybrid_forward(\n        self, F, x: Tensor, bin_centers: Tensor\n    ) -> Tuple[Tensor, Tensor]:\n        ps = self.proj(x)\n        reshaped_probs = ps.reshape(shape=(-2, -1, self.num_bins), reverse=1)\n        bin_centers = F.broadcast_add(bin_centers, ps.zeros_like())\n        return reshaped_probs, bin_centers\n\n\nclass BinnedOutput(DistributionOutput):\n    distr_cls: type = Binned\n\n    @validated()\n    def __init__(self, bin_centers: mx.nd.NDArray) -> None:\n        self.bin_centers = bin_centers\n        self.num_bins = self.bin_centers.shape[0]\n        assert len(self.bin_centers.shape) == 1\n\n    def get_args_proj(self, *args, **kwargs) -> gluon.nn.HybridBlock:\n        return BinnedArgs(self.num_bins, self.bin_centers)\n\n    def distribution(self, args, scale=None) -> Binned:\n        probs = args[0]\n        bin_centers = args[1]\n        F = getF(probs)\n\n        bin_centers = F.broadcast_mul(bin_centers, F.ones_like(probs))\n\n        if scale is not None:\n            bin_centers = F.broadcast_mul(\n                bin_centers, scale.expand_dims(axis=-1)\n            )\n\n        return Binned(probs, bin_centers)\n\n    @property\n    def event_shape(self) -> Tuple:\n        return ()\n", "path": "src/gluonts/distribution/binned.py"}]}