Repository: pyro-ppl/numpyro

Issue:
Usage of the `scale` effect handler?
The `scale` effect handler was introduced in commit [495cdbf2](https://github.com/pyro-ppl/numpyro/commit/495cdbf2c696ca9568e5a074e5b45d2e9eb2784c) but varies from other handlers in that it doesn't accept a callable `fn` argument on construction but only the scaling factor. Therefore I'm unsure how to use `scale` with any given model...

I'm suspecting that the `fn` argument should be there but has gone missing. Could it be that the implementation of `scale` has been taken more or less directly from Pyro's `ScaleMessenger` (where `Messenger` classes generally don't take callables in the constructor, however the `scale` convenience function does) and escaped the necessary adoptions for numpyro?


Discussion:
> Could it be that the implementation of scale has been taken more or less directly from Pyro's ScaleMessenger (where Messenger classes generally don't take callables in the constructor, however the scale convenience function does) and escaped the necessary adoptions for numpyro?

I think that's right, thanks for bringing this up. I think this was intended to be used as a context manager, but there is no reason why this shouldn't be a regular effect handler. I'll fix this shortly.

Resolving patch:
diff --git a/numpyro/handlers.py b/numpyro/handlers.py
--- a/numpyro/handlers.py
+++ b/numpyro/handlers.py
@@ -247,11 +247,11 @@ class scale(Messenger):
 
     :param float scale_factor: a positive scaling factor
     """
-    def __init__(self, scale_factor):
+    def __init__(self, fn=None, scale_factor=1.):
         if scale_factor <= 0:
             raise ValueError("scale factor should be a positive number.")
         self.scale = scale_factor
-        super(scale, self).__init__()
+        super(scale, self).__init__(fn)
 
     def process_message(self, msg):
         msg["scale"] = self.scale * msg.get('scale', 1)


Test patch:
diff --git a/test/test_handlers.py b/test/test_handlers.py
--- a/test/test_handlers.py
+++ b/test/test_handlers.py
@@ -1,3 +1,4 @@
+import pytest
 from numpy.testing import assert_allclose
 
 from jax import random
@@ -5,16 +6,20 @@
 import numpyro.distributions as dist
 from numpyro.handlers import sample, scale
 from numpyro.hmc_util import log_density
+from numpyro.util import optional
 
 
-def test_scale():
+@pytest.mark.parametrize('use_context_manager', [True, False])
+def test_scale(use_context_manager):
     def model(data):
         x = sample('x', dist.Normal(0, 1))
-        with scale(10):
+        with optional(use_context_manager, scale(scale_factor=10)):
             sample('obs', dist.Normal(x, 1), obs=data)
 
+    model = model if use_context_manager else scale(model, 10.)
     data = random.normal(random.PRNGKey(0), (3,))
     x = random.normal(random.PRNGKey(1))
     log_joint = log_density(model, (data,), {}, {'x': x})[0]
-    expected = dist.Normal(0, 1).log_prob(x) + 10 * dist.Normal(x, 1).log_prob(data).sum()
+    log_prob1, log_prob2 = dist.Normal(0, 1).log_prob(x), dist.Normal(x, 1).log_prob(data).sum()
+    expected = log_prob1 + 10 * log_prob2 if use_context_manager else 10 * (log_prob1 + log_prob2)
     assert_allclose(log_joint, expected)