openxla / xla

A machine learning compiler for GPUs, CPUs, and ML accelerators
Apache License 2.0
2.39k stars 356 forks source link

Cannot use `XlaOp` from outer scope #14299

Open joelberkeley opened 2 days ago

joelberkeley commented 2 days ago

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 commented 1 day ago

@GleasonK WDYT?