Skip to content

Conversation

@lexu27
Copy link
Collaborator

@lexu27 lexu27 commented Oct 19, 2025

This PR adds the following changes/fixes/enhancements:

  1. Support for non rank polymorphic tensor declaration using rankPrecondition :: RClassIdentifier -> Int -> DSLContext (). Users can now also use newSingletonRClass and newSingletonRClasses to create RClasses with a fixed rank of 1. The bound inference algorithm sets upper and lower bounds for fixed rclasses equal to the number the users set. iota and concatTensor call markSingleton under the hood.

  2. printTitle prints with colors for whether verification succeds or fails.

  3. 11/12 TASO operators are implemented now excluding pool. The operators are desugared down to XLA operators in TASO.hs.

  4. Most operators from the TASO paper are verified within TensorRight (except for pool and a few rules that either cannot be verified or hang). Only matmul and conv have troublesome properties.

Problematic Rules

  1. $∀x \ matmul(x,\ I_{matmul}) = x$

    • When doing doing matmul with $x$ and $I$ it's not clear how to create an SI-Relation between the contraction axis for the matmul and $x$.
  2. $∀x, y, z, w \ matmul(concat(1, x, z),\ concat(0, y, w)) = ewadd(matmul(x,\ y),\ matmul(z,\ w))$

    • This rule does not currently pass. It may be because you have to split the lhs contraction axis into piecewise parts to match the 2 contraction axes on the rhs.
  3. Any convolution rule with SAME padding currently needs more than 15 seconds to verify so the rewrite rules time out. This is because of the complexity with dynamically calculating the proper padding configuration for SAME padding based on input and kernel maps.

  4. $∀k, x \ conv(1, P_{same}, A_{none}, x, I_{conv}(k)) = x$

    • convBaseCheck forces output feature classes to not be in the inputRClasses. Since we are comparing to x this might cause trouble during verification as the output feature rClass is no longer the same as the input.
  5. $∀s, p, x, y, z, w \ conv(s, p, A_{none}, concat(1, x, z), concat(1, y, w)) = ewadd(conv(s, p, A_{none}, x, y), conv(s, p, A_{none}, z, w))$

    • Similar to the matmul rule above, we have to split the lhs contraction axis into two separate parts to match the contraction axes for each of the convolutions on the rhs.

Checks

  • Did you run ormolu, a Haskell Formatter? Run ormolu --mode inplace $(git ls-files '*.hs') to format all files in the project.
  • Did you update hie.yaml using the command gen-hie > hie.yaml?
  • Do all tests pass using stack test?

Not all TASO tests pass using stack test.

@lexu27 lexu27 requested a review from jaiarora0011 October 19, 2025 05:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants