Searched refs:iotas_to_concat (Results 1 – 2 of 2) sorted by relevance
116 std::vector<xla::XlaOp> iotas_to_concat; in Compile() local117 iotas_to_concat.push_back(i); in Compile()118 iotas_to_concat.push_back(j); in Compile()119 idx = xla::ConcatInDim(ctx->builder(), iotas_to_concat, 1); in Compile()
270 std::vector<XlaOp> iotas_to_concat; in CompileWhereWithPrefixSum() local274 iotas_to_concat.push_back( in CompileWhereWithPrefixSum()277 XlaOp iotas = xla::ConcatInDim(b, iotas_to_concat, /*dimension=*/1); in CompileWhereWithPrefixSum()