Test for whether TransformModules work for density estimation

[stefanwebb]

I've added a test for whether TransformModule's work for density estimation after discovering that Spline doesn't... I'll fix the bug in Spline in this PR

[stefanwebb]

I worked out that the bug was that several transforms didn't update their parameters for each call of the forward or inverse operation. This is a problem after you've taken a gradient step during learning...

We now test that density estimation is possible and it passes for all transforms!

[stefanwebb]

@fritzo when this merges it will unblock #2542

[fritzo]

Not part of this PR, but it is an antipattern to make return type depend on a scalar value. Doing so breaks the ability to write generic code using this interface. Has this behavior been released yet?

[stefanwebb]

Yes I'm afraid it's been released... Can you elaborate on why this would break the ability to write generic code? I think the underlying problem is actually how DenseNN shapes it's outputs, and this is something we should think about more carefully during a major refactoring

[fritzo]

warning: the following is highly subjective

An example of some generic code you might like to write is, say, a ConditionalCatTransform that concatenates two conditional transforms. You might want to fuse their neural nets to save on tensor ops and share strength. How would you do this? Maybe

class ConditionalCatTransform(...):
    def __init__(self, parts):
        super().__init__(self, cache_size=1)
        self.parts = parts  # ignoring ModuleList magic for demo purposes
        self.nn = memoize(cat_dense_nn([part.nn for part in self.parts]))
        end = 0
        for part in self.parts:
            beg, end = end, end + part.count_transforms
            # The following line assumes DenseNN returns a tuple:
            part.nn = lambda context: self.nn(context)[beg:end]

then we could define helpers memoize = functools.lru_cache(max_size=1) and

def cat_dense_nn(parts):
    input_dims = parts[0].input_dims
    hidden_dims = parts[0].hidden_dims
    param_dims = sum([part.param_dims for part in parts], [])
    return DenseNN(input_dims, hidden_dims, param_dims) 

That's what we could have written if the types were consistent. It would have been pretty simple generic code. But it looks like DenseNN returns an output type that depends on an int value, so we would need to complicate our wrapper with extra logic:

  class ConditionalCatTransform(...):
      def __init__(self, parts):
          ...
          for part in self.parts:
              beg, end = end, end + part.count_transforms
-             part.nn = lambda context: self.nn(context)[beg:end]
+             if beg + 1 == end:
+                 part.nn = lamba context: self.nn(context)[beg]
+             else:
+                 part.nn = lambda context: self.nn(context)[beg:end]

Now that's just a little more complex. But I feel, as an author of abstract code, that the complexity tax is best paid by one-off code, so that abstractions can be built tax free. I think this tax structure is one of the main differences between programming languages, e.g. R and MATLAB tend to tax library code, whereas Python and C tend to tax one-off code.

[stefanwebb]

Okay, let's think about this when we have a chance to overhaul DenseNN

[stefanwebb]

@fritzo I'm not sure what's wrong with the AIR tutorial in Travis...

[fritzo]

@neerajprad Do you have any ideas why the air tutorial might be failing? I haven't been able to reproduce locally.

[neerajprad]

I think #2549 should temporarily silence the FutureWarning from jupyter_client that's causing the tutorial tests to fail.

[fritzo]

@stefanwebb can you try merging in dev and pushing to trigger ci?

[stefanwebb]

Thanks @neerajprad! It all passes now 😄