// Add inline pass.
  pass_manager.addPass(mlir::createInlinerPass());

  // Expands mhlo.tuple ops.
  pass_manager.addPass(
      mlir::mhlo::createExpandHloTuplesPass(entry_function_name.str()));
  // Flatten tuples for control flows.
  pass_manager.addNestedPass<mlir::func::FuncOp>(
      mlir::mhlo::createFlattenTuplePass());

  mlir::odml::AddMhloOptimizationPasses(
      pass_manager,

    Value n_plus_y = rewriter.create<AddOp>(loc, iotaN, y);

    // GatherOp is happy about letting us index out of bounds values, but those
    // values will be undefined. So we mask them later. Set up the boolean
    // expression that tells us which entries, in the output shape, are out of
    // bounds and thus become the padding_value.
    Value x_in_bounds = rewriter.create<AndOp>(
        loc,