Cannot use XlaOp from outer scope

[joelberkeley]

In the following MWE

        xla::XlaBuilder root("root");
        auto zero = xla::ConstantR0(&root, 0);
        absl::Span<const int64_t> xs = {0};
        auto zeros = xla::ConstantR1(&root, xs);

        auto add_builder = root.CreateSubBuilder("add");
        auto u32 = xla::ShapeUtil::MakeScalarShape(xla::PrimitiveType::U32);
        auto p0 = xla::Parameter(&*add_builder, 0, u32, "p0");
        auto p1 = xla::Parameter(&*add_builder, 1, u32, "p1");
        auto add = xla::Add(xla::Add(p0, p1), zero);
        auto add_comp = add_builder->Build().value();

        auto res = xla::Reduce(zeros, zero, add_comp, {0});
        root.Build().value();

I see the error

terminate called after throwing an instance of 'absl::lts_20230802::BadStatusOrAccess'
what(): Bad StatusOr access: INVALID_ARGUMENT: XlaOp with handle 1 is built by builder 'root', but is trying to use it in builder 'add':

which suggests to me that it's not possible to use an XlaOp from an outer scope (XlaBuilder) in an inner scope. Is this a bug? Note that in a more complex example (double nested Map) I get a different error that points to the same problem

[external/com_google_absl/absl/status/statusor.cc : 83] RAW: Attempting to fetch value instead of handling error INVALID_ARGUMENT: XlaOp with handle 2 is built by builder 'root:map/f', but is trying to use it in builder 'root:map/f:map/f'

[cheshire]

@GleasonK WDYT?

[GleasonK]

I believe this is intended behavior for HLO. While this is supported in MLIR, it hasn't been supported in HLO -- I can't speak to whether this is a compiler invariant or just a feature that was never explored/added to HLO.

We have prepare_for_export.cc > prepareExplicitCapturedConstants which runs before MHLO->HLO conversion to ensure all values are explicitly captured in the proper scope.

[joelberkeley]

ok thanks. I'm good if it's possible in StableHLO. I'm guessing it is if it's supported in MLIR. If this is intended, we can close the ticket and I'll just warn users until I migrate to StableHLO