Leveraging DenseArray attributes in StableHLO

[Michael Levesque-Dion]

Hi everyone,

If you do not have a dependency on StableHLO and/or are not interested in StableHLO implementation details, you can stop reading.

This post is a heads up about a change we are making to op fields representing 1-dimensional arrays. Until recently, these fields used I64ElementsAttr (example), which actually allows tensors of any dimension. In MLIR source, this typically looks like e.g. dense<1, 2, 3> : tensor<3xi64>. We are currently in the process of migrating such fields to instead use DenseI64ArrayAttr (example), which is a better fit for what these fields are trying to represent. In MLIR source, this typically looks like e.g. array<i64: 1, 2, 3>.

We have already migrated some fields to use DenseI64ArrayAttr, but more recently we introduced an attribute that can handle both dense elements and dense arrays to ease the migration. Concretely, this means you can update your code to use arrays now and avoid issues when we actually migrate the fields to DenseI64ArrayAttr. See the list of affected fields.

We are planning to start migrating everything to DenseI64ArrayAttr soon. We don't expect any interruptions to existing workflows. If you run into any issues, most likely the cause is version skew between frameworks and XLA dependencies. Feel free to reach out if you run into any issues. Updates will be provided via an issue on GitHub.

Special thanks to Mehdi for getting this work started!

Cheers,

Michael

Questions & Answers

Q1: Will this have any impact on VHLO forward/backward compatibility?
A1: There will be no impact on VHLO forward/backward compatibility. The same VHLO will be produced before and after the change.

Q2: What about MHLO?
A2: Performing this migration for StableHLO has proven to be fairly time consuming, so we are still figuring out a plan for migrating MHLO.