}

inline bool IsJaxRandomUniform(mlir::func::FuncOp func) {
  return func.getName().contains("tfl_wrapped_jax_random_uniform");
}

inline bool IsJaxRandomNormal(mlir::func::FuncOp func) {
  return func.getName().contains("tfl_wrapped_jax_random_normal");
}

void LegalizeJaxRandomPass::runOnOperation() {
  auto func = getOperation();
  if (!IsJaxRandomUniform(func) && !IsJaxRandomNormal(func)) return;