// CHECK-NEXT: "tf.AssignVariableOp"(%[[VAR]], %[[IDENTITY]])
}

// -----

// If something other than `tf.VarHandleOp` is being initialized, it is
// not erased.

// CHECK-LABEL: module
module attributes {tf_saved_model.semantics} {
  "tf_saved_model.session_initializer"() {initializers = [@init_func_restore_op]} : () -> ()
  // CHECK: "tf_saved_model.session_initializer"()

        op_builder->create<mhlo::BroadcastInDimOp>(
            const_op->getLoc(), splat_elements_attr.getType(), scalar,
            broadcast_dims);
    const_op->replaceAllUsesWith(broadcast_in_dim_op);
    const_op->erase();
  }
};

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateUnfoldSplatConstantPass() {
  return std::make_unique<UnfoldSplatConstantPass>();
}

                         op->getAttrs(), op->getSuccessors());
    Operation* inlined = builder.create(state);

    parent_op->replaceAllUsesWith(inlined);
    parent_op->erase();
  }
  return success();
}

// Three instances of the rule to cover the three different types of
// TF::FakeQuant operators
using PreparePerTensorFakeQuant = InsertTFLQuantOpsAfterTFFakeQuantOp<

    auto index = indexed_assign.index();
    assign_op->moveAfter(new_replicate_op);
    assign_op->setOperand(
        1, new_replicate_op->getResult(old_num_results + num_replicas * index));
  }
  replicate_op->erase();
}

// Looks for AssignVariable ops from the end of the tf_device.replicate op. It
// returns all the last writes to replicate invariant resource variables

  CreateRegionWithCall(while_op.cond_function(), new_op.getCond(), loc);
  CreateRegionWithCall(while_op.body_function(), new_op.getBody(), loc);

  op->replaceAllUsesWith(new_op.getResults());
  op->erase();
}

void LegalizeWhilePass::RunOnFunction(func::FuncOp func) {
  // Convert all TF WhileOps inside the function body to TFL While ops.
  func.getBody().walk([](TF::WhileOp while_op) { RunOnWhile(while_op); });
}

      ValueRange(resource_args), cluster_func_op.getFuncAttr());

  CopyDeviceAndUnderscoredAttributes(cluster_func_op, xla_launch_op);
  cluster_func_op.replaceAllUsesWith(xla_launch_op.getResults());
  cluster_func_op.erase();
}

void XlaRewritePass::runOnOperation() {
  ModuleOp module = getOperation();
  SymbolTable symtab(module);
  OpBuilder builder(&getContext());
  module.walk([&](tf_device::ClusterFuncOp cluster_func_op) {

    DenseElementsAttr attr =
        DenseElementsAttr::get(type, b.getI32IntegerAttr(index));
    auto constant = b.create<arith::ConstantOp>(op.getLoc(), type, attr);
    op.replaceAllUsesWith(constant.getOperation());
    op.erase();
  });
}

// Creates an instance of the TensorFlow DeviceIndex selector pass.
std::unique_ptr<OperationPass<func::FuncOp>> CreateDeviceIndexSelectorPass() {
  return std::make_unique<DeviceIndexSelector>();
}

            init_op.getLoc(), init_op.getTableHandle(), const_op.getResult(),
            init_op.getKeyIndex(), init_op.getValueIndex(),
            init_op.getVocabSize(), init_op.getDelimiter());
        init_op.erase();
      }
    }
    func.eraseArguments(args_to_erase);
  }
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateFreezeAssetsPass(
    std::string saved_model_dir) {

    Operation &return_op = block.front();
    for (auto ret_mapping :
         llvm::zip(op->getResults(), return_op.getOperands())) {
      std::get<0>(ret_mapping).replaceAllUsesWith(std::get<1>(ret_mapping));
    }
    op->erase();
  });
}

}  // namespace

namespace TF {
std::unique_ptr<OperationPass<func::FuncOp>>
CreateMaterializePassthroughOpPass() {
  return std::make_unique<MaterializePassthroughOpPass>();
}

    TFL::FillOp fill = op_builder->create<TFL::FillOp>(
        const_op->getLoc(), splat_elements_attr.getType(), fill_shape,
        fill_value);
    const_op->replaceAllUsesWith(fill);
    const_op->erase();
  }
};

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateUnfoldLargeSplatConstantPass() {
  return std::make_unique<UnfoldLargeSplatConstantPass>();
}

}  // namespace TFL