Implement plated einsum#692
Conversation
eb8680
left a comment
eb8680
left a comment
There was a problem hiding this comment.
Getting close! Seems like mask has clarified things.
| return fwd() | ||
|
|
||
|
|
||
| class ReduceDeltaSimpleRange(ObjectInterpretation): |
There was a problem hiding this comment.
Can we decompose ReduceDeltaSimpleRange into an equational rewrite step in ops.monoid and simpler kernel invocations in handlers.jax.monoid? I think this and ReduceDependentRangeMask are the last rules in jax.monoid that do any significant non-JAX-specific rewriting.
| return fwd() | ||
|
|
||
|
|
||
| # Cross-cutting delta rules not yet implemented: |
There was a problem hiding this comment.
Is this comment obsolete with the new mask semantics? Either way it should probably be updated and moved or removed so it doesn't go stale.
| return bind_dims(inner, fresh_op) | ||
|
|
||
|
|
||
| class ReduceDependentRangeMask(ObjectInterpretation): |
There was a problem hiding this comment.
Can we decompose ReduceDependentRangeMask into an equational rewrite step in ops.monoid and simpler kernel invocations in handlers.jax.monoid? Is this even still necessary following the move to mask? I think this and ReduceDeltaSimpleRange are the last rules in jax.monoid that do any significant non-JAX-specific rewriting.
| peeled_cprod_streams = { | ||
| k: v for (k, v) in cprod_streams.items() if k != plate_op | ||
| } | ||
| if not peeled_cprod_streams: |
There was a problem hiding this comment.
nit: let the generic reduce(body, {}) rule deal with this conditional instead of inlining it, in the interest of slightly simplifying an already quite complex rule body.
| inner_tail_streams = { | ||
| k: v for (k, v) in inner_streams.items() if k != inner_plate_op | ||
| } | ||
| if inner_tail_streams: |
There was a problem hiding this comment.
nit: let the generic reduce(body, {}) rule deal with this conditional instead of inlining it, in the interest of slightly simplifying an already quite complex rule body.
|
|
||
| # include any extra sum streams outermost | ||
| tail_streams = {k: v for (k, v) in streams.items() if k != stream_key} | ||
| if tail_streams: |
There was a problem hiding this comment.
nit: let the generic reduce(body, {}) rule deal with this conditional instead of inlining it, in the interest of slightly simplifying an already quite complex rule body.
| monoid.reduce(group(using), inner_streams), | ||
| monoid.reduce(group(rest), inner_streams), | ||
| ) | ||
| return monoid.reduce(body2, outer_streams) if outer_streams else body2 |
There was a problem hiding this comment.
nit: let the generic reduce(body, {}) rule deal with this conditional instead of inlining it, in the interest of slightly simplifying an already quite complex rule body.
| continue | ||
|
|
||
| peeled_body = Union.reduce( | ||
| [Union.delta(drop_elem(idx, plate_index), union_body)], union_streams |
There was a problem hiding this comment.
It looks like this is the only place in ops.monoid where a fresh delta term is introduced. Is there a way to express the rule in terms of mask instead? It might be helpful for limiting the complexity of the normalization rules to avoid introducing deltas ourselves, even if we keep it in the Monoid interface.
|
|
||
| for i, elem in enumerate(mask_elems): | ||
| match elem: | ||
| case Term(_NumberTerm.__ne__, (Term(stream_op, (), {}), index)) if ( |
There was a problem hiding this comment.
Should this part of ReduceArrayScan be a separate rule in ops.monoid? The next case is array-specific, but this one isn't.
2 participants
No description provided.