This has two problems:
* it's not thread-safe, which will become problematic if we run tests with thread-parallelism.
* it's not very maintainable.
Instead, add a new util.test_event(...) function that can be called at points of interest in the program. test_utils registers a callback that is invoked when an event is received. This avoids the need to make thread-unsafe global monkey patches.
These constraints turn out to be quite useful, e.g., when
we want to say that certain dimensions are a multiple of
a device axis.
Previously, the constraint `mod(e, k) == 0` was being useful
only to normalize away `mod(e, k)`. In particular it was not
useful for proving `k * floordiv(e, k)`. Now we add that
features.
The values are guaranteed to be 0 or 1 since we create this array ourselves when processing the masks into a MaskInfo object.
PiperOrigin-RevId: 705252534
Layouts are added as annotations on MLIR ops, using the `in_layouts` and
`out_layouts` attributes.
At this point, layout inference is done in two passes: one "backwards" pass
(root-to-parameters), and one "forward" pass (parameters-to-root).
Each pass goes through all the ops in the specified order, and infers a
possible layout from the layout information that is available. We expect to
need two passes because partial layout annotations may be provided on
intermediate nodes (e.g. `wgmma`), and a single pass from the root to the
parameters is therefore insufficient to properly annotate all the operations.
We do not perform any check as to whether the inferred layouts can be further
lowered correctly---meaning that the produced IR can possibly fail to lower
later.
Layouts are only inferred for ops involving at least one operand or result of
type `VectorType`/`RankedTensorType`.
When layouts can't be inferred for an op that should have them, we default to
annotating it with strided fragmented layouts.
PiperOrigin-RevId: 705092403
Without this type hint, some tools (including PyCharm) infer the more
generic return type from typing.NamedTuple.
To improve user experience, I've added a narrower type hint.
However, the typing of this method is still 'flawed' as the only properly supported
input is another Rotation. This is a narrower input type and therefore
violates the Liskov substitution principle. Therefore I left the input
parameter untyped.
For more info:
https://mypy.readthedocs.io/en/stable/common_issues.html#incompatible-overrides
Previously, an equality constraint was used only as a normalization
rule. This created a problem for constraints of the form "4*b=c",
because it would not allow proving that "b <= c" (since the
normalization of "4*b" kicks in only if "b" is multiplied by a
multiple of 4.
Now we add the equality constraints also in the inequality
reasoning state.
