if (failed(layout)) return failure();
      v.push_back(layout.value());
    }
    ArrayRef<Attribute> shape(v);
    return rewriter.getArrayAttr(shape);
  } else if (mlir::isa<TupleType>(types[0])) {
    auto tuple_type = mlir::dyn_cast<TupleType>(types[0]);
    const auto &types = tuple_type.getTypes();
    llvm::SmallVector<mlir::Attribute> v;
    v.reserve(types.size());
    for (const mlir::Type &t : types) {

}

void LegalizeJaxRandomPass::runOnOperation() {
  auto func = getOperation();
  if (!IsJaxRandomUniform(func) && !IsJaxRandomNormal(func)) return;
  auto result_tuple_ty =
      mlir::dyn_cast_or_null<TupleType>(func.getFunctionType().getResult(0));
  if (!result_tuple_ty) return;
  if (result_tuple_ty.size() != 1) return;
  auto result_ty = mlir::dyn_cast<ShapedType>(result_tuple_ty.getType(0));

  func.eraseBody();

    auto new_result = CreateTuple(builder, value.getLoc(), operands);
    rewritten_values.insert({value, new_result});
    return new_result;
  };

  auto tuple_type = mlir::dyn_cast<TupleType>(value.getType());
  // `value` is not a tuple, create a new tuple.
  if (!tuple_type) return {create_tuple({value, token})};

  // Extend tuple if `value` is a tuple.

    if (while_op->getNumRegions() != 2) return failure();

    // This rule doesn't support mhlo::WhileOp with tuple inputs.
    for (auto type : while_op->getOperandTypes()) {
      if (mlir::isa<TupleType>(type)) return failure();
    }

    // Creates a TF::WhileRegionOp to replace the mhlo::WhileOp. HLO WhileOp
    // currently doesn't support stateless and shape invariant, so these