Allow shapecheck of PixelCNN++

[juliuskunze]

Makes shapecheck work on PixelCNN++. Allows shapecheck of indexing, slicing, broadcast_to, reshape. Updated: np.split, np.where and random sampling are now taken out of this PR for simplification. They required special handling for polymorphic shapes since

@shapecheck(['n'], 'n')
  def neg_zeros_like(x):
    return -np.zeros(x.shape)

breaks: find_top_trace in full_p.bind returns None here since all arguments are independent of any inputs. Therefore, the result of zeros is not wrapped in a trace, and the np.negative implementation is called with a polymorphic shape instead of the shape rule, causing an error. This also breaks shape check of -np.ones(x.shape), -np.broadcast_to(0, x.shape), -np.arange(10) etc.

I will keep the workarounds in a separate branch until this issue is fixed.

[mattjj]

One way to think about this is that np.zeros, np.ones, np.arange etc take no value-level arguments and only take type-level arguments. In NumPy there isn't much of a distinction, but in JAX land we think of shapes as part of the type system.

But as you've seen our np.zeros implementation expects to get a tuple of integers, not these new parametrically-polymorphic type arguments that masking.py is effectively adding to the system. This also interacts with our tracing mechanism, as you pointed out, which is based on data dependence of values: here np.zeros(x.shape) has no data dependence on a value, yet it needs to produce a value with a polymorphic shape (i.e. lifted into the appropriate ShapeCheckTrace instance).

This is tricky to solve, but doable. We had something analogous with pmap and things like lax.psum(1, 'i') or lax.axis_index('i'), where again there's no data dependence on a traced value. Instead we use the fact that the axis name 'i' corresponds to a trace, and we use that together with a global mapping from axis names to traces (dynamic_axis_env in pxla.py) to lift into the trace.

I think we need to do something similar here: every shape indeterminate is associated with a trace; that information isn't explicitly stored anywhere but you could imagine adapting shape_envs to include this mapping. Then when we call np.zeros and notice that we get a polymorphic shape argument, we can use that mapping to lift into the appropriate trace.

Pretty tricky! Maybe we should pair program it :)

WDYT? (I could be wrong about any part of the above!)

[juliuskunze]

@mattjj That's a very neat approach and might just be what we need. In our case, we could simply attach a ShapeCheckTrace to Poly. This should always be possible since they are currently only created

• from ShapeSpecs, which are passed as arguments into shapecheck, where we can access the corresponding ShapeCheckTrace, or
• from existing Polys via operators.

I haven't looked into the pmap implementation yet, but I would most certainly need help with the unusual lifting. I guess pair programming might make sense here. (:

Also, it looks like some tests break because onp.float64 values are used as indices or sizes. This is only now an issue since I replaced some int(x) calls with x.__index__() to allow for polymorphic shapes, resulting in stricter type requirements. Update x2: This is now fixed, and came from the surprising behavior that 1 * onp.array(1, onp.uint64) returns a float64 here (in contrast to 1 * onp.array(1, onp.uint32) returning an int as expected). Relatedly, lax.iota is now stricter: size can't be a float anymore, conforming with the XLA spec. (np.arange is unaffected.)

[juliuskunze]

@mattjj We might want to break up the tests into 2 files (shapecheck and masking), maybe even the interpreter file. Thoughts?

[mattjj]

Splitting files SGTM, whatever you think is best.

Re attaching traces to Polys, we might have a lifetime problem: Polys can outlive the trace, embedded in jaxprs, but we don’t want to leak trace objects. We could fix that with weakrefs but having an explicitly managed data structure might be better.

[juliuskunze]

@j-towns and I found that shapecheck does not require a custom trace(r). Instead, we can reuse abstract_eval_fun to propagate polymorphic shapes. This works since polymorphic shape evaluation is now based on the default shape rules. For this reason, I removed ShapeCheckTrace(r) and the shape_rules dictionary.

I also moved all remaining code needed by shapecheck out of masking.py:

• I moved Poly and related code (~160 lines) into abstract_arrays.py, since any code handling abstract arrays should also be able to handle polymorphic shapes now.
• The shape spec parser (~70 lines) is now in api.py, where it is used for shapecheck and mask.

@mattjj @shoyer Let me know if you know a better location for these.

[mattjj]

Overall that reorganization sounds reasonable, but you need to split this up into smaller changes to make it easier to review. For example, one PR could just be moving code without adding any features / methods etc. Another could be deleting the shape_rules / ShapeCheckTracer stuff. Other PRs could be adding specific features. (Also, none of the PRs should make unrelated changes, like rewriting import statements for unrelated code, or changing existing functions to use safe_map instead of map, etc.)

WDYT?

[mattjj]

Thanks for continuing to iterate and simplify!

Still I think we should split this into multiple minimal PRs, where by minimal I mean smallest change unit. That makes it easier to review (reducing the probability we introduce bugs), easier to find the provenance of a bug later (using git bisect) and easier to roll any parts back if we need. This is something I’m trying to get better at myself :)

WDYT? What are the minimal units here?

[juliuskunze]

@mattjj As discussed, I reverted the relocation for now to simplify the diff. I will create a separate PR with these changes once this one is merged.

[juliuskunze]

@mattjj I got tons of improvements in the pipeline:

• I got the custom lifting+abstractifying for polymorphic shapes working, solving the problem discussed above. This allows shapecheck of zeros, broadcast_to, arange, split, where, random.uniform without special case workarounds.
• mask of jit, conv, gather, pad...

They all depend on this PR + the relocation, so I'll have to wait with PRs until this is merged. FYI I have time to work on merging this back until Friday, can't promise anything afterwards.

[juliuskunze]

@mattjj I got ragged batching on PixelCNN++ (vmap(mask(pixelcnn))) to work. I can prepare a PR as soon as this is merged. Is there anything holding back a merge?

[mattjj]

Nothing holding us back other than me being slow! Thanks for your patience, and for this amazing contribution.

[mattjj]

Woohoo, so much progress!

[gnecula]

I am reverting this commit until we figure out how to handle this (a user complained, and there may be some other test failures).