iree: Failure to lower mhlo.tuple in Brax program

I’m trying to compile a program coming from Brax with this Colab notebook.

Reproducers here: https://gist.github.com/ScottTodd/4f4434a25f3036f1a37cda88aa872add

Compilation from Python (with asserts disabled?) fails with

error: failed to legalize operation 'mhlo.dot' that was explicitly marked illegal

Repro command using native tools with asserts enabled: iree-translate -iree-input-type=xla -iree-mlir-to-vm-bytecode-module -iree-hal-target-backends=dylib-llvm-aot inference_fn_tf-mhlo.mlir

Callstack with the assert:

Assertion failed: Index < Length && "Invalid index!", file D:\dev\projects\iree\third_party\llvm-project\llvm\include\llvm/ADT/ArrayRef.h, line 255

iree-translate.exe!HandleAbort(int Sig) Line 408 (d:\dev\projects\iree\third_party\llvm-project\llvm\lib\Support\Windows\Signals.inc:408)
ucrtbase.dll!00007fff4ba31881() (Unknown Source:0)
ucrtbase.dll!00007fff4ba32851() (Unknown Source:0)
ucrtbase.dll!00007fff4ba341b5() (Unknown Source:0)
ucrtbase.dll!00007fff4ba344f1() (Unknown Source:0)
[Inline Frame] iree-translate.exe!llvm::ArrayRef<__int64>::operator[](unsigned __int64) Line 255 (d:\dev\projects\iree\third_party\llvm-project\llvm\include\llvm\ADT\ArrayRef.h:255)
iree-translate.exe!mlir::`anonymous namespace'::GetDotOperationType(mlir::mhlo::DotOp dot_op) Line 1371 (d:\dev\projects\iree\third_party\mlir-hlo\lib\Dialect\mhlo\transforms\legalize_to_linalg.cc:1371)
iree-translate.exe!mlir::`anonymous namespace'::DotOpConversion<1,mlir::linalg::MatvecOp>::matchAndRewrite(mlir::mhlo::DotOp op, mlir::mhlo::DotOpAdaptor adaptor, mlir::ConversionPatternRewriter & rewriter) Line 1414 (d:\dev\projects\iree\third_party\mlir-hlo\lib\Dialect\mhlo\transforms\legalize_to_linalg.cc:1414)
iree-translate.exe!mlir::OpConversionPattern<mlir::mhlo::DotOp>::matchAndRewrite(mlir::Operation * op, llvm::ArrayRef<mlir::Value> operands, mlir::ConversionPatternRewriter & rewriter) Line 427 (d:\dev\projects\iree\third_party\llvm-project\mlir\include\mlir\Transforms\DialectConversion.h:427)
iree-translate.exe!mlir::ConversionPattern::matchAndRewrite(mlir::Operation * op, mlir::PatternRewriter & rewriter) Line 1702 (d:\dev\projects\iree\third_party\llvm-project\mlir\lib\Transforms\Utils\DialectConversion.cpp:1702)
iree-translate.exe!mlir::PatternApplicator::matchAndRewrite(mlir::Operation * op, mlir::PatternRewriter & rewriter, llvm::function_ref<bool __cdecl(mlir::Pattern const &)> canApply, llvm::function_ref<void __cdecl(mlir::Pattern const &)> onFailure, llvm::function_ref<mlir::LogicalResult __cdecl(mlir::Pattern const &)> onSuccess) Line 201 (d:\dev\projects\iree\third_party\llvm-project\mlir\lib\Rewrite\PatternApplicator.cpp:201)
iree-translate.exe!`anonymous namespace'::OperationLegalizer::legalizeWithPattern(mlir::Operation * op, mlir::ConversionPatternRewriter & rewriter) Line 1985 (d:\dev\projects\iree\third_party\llvm-project\mlir\lib\Transforms\Utils\DialectConversion.cpp:1985)
iree-translate.exe!`anonymous namespace'::OperationLegalizer::legalize(mlir::Operation * op, mlir::ConversionPatternRewriter & rewriter) Line 1891 (d:\dev\projects\iree\third_party\llvm-project\mlir\lib\Transforms\Utils\DialectConversion.cpp:1891)
iree-translate.exe!`anonymous namespace'::OperationConverter::convert(mlir::ConversionPatternRewriter & rewriter, mlir::Operation * op) Line 2404 (d:\dev\projects\iree\third_party\llvm-project\mlir\lib\Transforms\Utils\DialectConversion.cpp:2404)
iree-translate.exe!`anonymous namespace'::OperationConverter::convertOperations(llvm::ArrayRef<mlir::Operation *> ops, llvm::function_ref<void __cdecl(mlir::Diagnostic &)> notifyCallback) Line 2453 (d:\dev\projects\iree\third_party\llvm-project\mlir\lib\Transforms\Utils\DialectConversion.cpp:2453)
iree-translate.exe!mlir::applyPartialConversion(llvm::ArrayRef<mlir::Operation *> ops, mlir::ConversionTarget & target, const mlir::FrozenRewritePatternSet & patterns, llvm::DenseSet<mlir::Operation *,llvm::DenseMapInfo<mlir::Operation *,void> > * unconvertedOps) Line 3267 (d:\dev\projects\iree\third_party\llvm-project\mlir\lib\Transforms\Utils\DialectConversion.cpp:3267)
iree-translate.exe!mlir::applyPartialConversion(mlir::Operation * op, mlir::ConversionTarget & target, const mlir::FrozenRewritePatternSet & patterns, llvm::DenseSet<mlir::Operation *,llvm::DenseMapInfo<mlir::Operation *,void> > * unconvertedOps) Line 3275 (d:\dev\projects\iree\third_party\llvm-project\mlir\lib\Transforms\Utils\DialectConversion.cpp:3275)
iree-translate.exe!mlir::iree_compiler::MHLO::`anonymous namespace'::ConvertMHLOToLinalgOnTensorsPass::runOnOperation() Line 367 (d:\dev\projects\iree\iree\compiler\InputConversion\MHLO\MHLOToLinalgOnTensors.cpp:367)
iree-translate.exe!mlir::detail::OpToOpPassAdaptor::run(mlir::Pass * pass, mlir::Operation * op, mlir::AnalysisManager am, bool verifyPasses, unsigned int parentInitGeneration) Line 393 (d:\dev\projects\iree\third_party\llvm-project\mlir\lib\Pass\Pass.cpp:393)
iree-translate.exe!mlir::detail::OpToOpPassAdaptor::runPipeline(llvm::iterator_range<llvm::pointee_iterator<std::unique_ptr<mlir::Pass,std::default_delete<mlir::Pass> > *,mlir::Pass> > passes, mlir::Operation * op, mlir::AnalysisManager am, bool verifyPasses, unsigned int parentInitGeneration, mlir::PassInstrumentor * instrumentor, const mlir::PassInstrumentation::PipelineParentInfo * parentInfo) Line 452 (d:\dev\projects\iree\third_party\llvm-project\mlir\lib\Pass\Pass.cpp:452)

this is the line from the top of the stack: https://github.com/tensorflow/mlir-hlo/blob/96fe02f272c4cf8c874a20497a902b59527ecab8/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc#L1371