auto collectOps = [this](Operation *op) {
    const auto &op_name = op->getName().getStringRef().str();
    if (absl::StartsWith(op_name, "tf.") || absl::StartsWith(op_name, "tfl.")) {
      loc_to_name_.emplace(op->getLoc(), op_name);
    }
  };

  for (auto &region : module->getRegions()) {
    region.walk(collectOps);
  }

  pass_name_ = extract_pass_name(pass->getName().str());
  error_collector_->Clear();
}

    for (int i = 0; i < op->getNumOperands(); i++) {
      Value input = op->getOperand(i);
      if (getElementTypeOrSelf(input).isBF16()) {
        Value f32_cast = rewriter.create<TF::CastOp>(
            op->getLoc(),
            CloneTypeWithNewElementType(input.getType(), rewriter.getF32Type()),
            input);
        op->setOperand(i, f32_cast);
      }
    }

    // Casts BF16 outputs of the operation.

    auto new_reshape_type = RankedTensorType::get(new_reshape_dims, el_ty);
    ReshapeOp new_reshape =
        rewriter.create<ReshapeOp>(new_reshape_shape.getLoc(), new_reshape_type,
                                   op.getInput(), new_reshape_shape);
    TF::ConstOp new_broadcast_shape =
        GetI64ConstantTensor(rewriter, reshape_shape, op.getLoc());

  for (Value result : new_launch_op_results)
    new_launch_op_results_types.push_back(result.getType());

  auto new_launch_op = builder.create<tf_device::LaunchOp>(
      old_launch_op->getLoc(), old_launch_op->getDeviceAttr(),
      /*result_types=*/new_launch_op_results_types);

  new_launch_op.getBody().takeBody(old_launch_op->getBody());

def CreateTFReadVariableOpFromResourceHandle : NativeCodeCall<
    "$_builder.create<TF::ReadVariableOp>("
    "$0.getLoc(), GetResourceSubtype($1), $1)">;

def CreateTFSelectOp: NativeCodeCall<
    "$_builder.create<TF::SelectOp>("
    "$0.getLoc(), $3.getType(), $1, $2, $3)">;

def ConstAttrIfThenElse: NativeCodeCall<
    "($0).getValue() ? ($1) : ($2)">;

  // When creating the new const op, use a fused location that combines the
  // original const and the qbarrier that led to the quantization.
  auto fusedLoc = rewriter.getFusedLoc(
      {qbarrier.getArg().getDefiningOp()->getLoc(), qbarrier.getLoc()});
  auto newConstOp = rewriter.create<arith::ConstantOp>(
      fusedLoc, newConstValueType, cast<TypedAttr>(newConstValue));
  rewriter.replaceOpWithNewOp<StorageCastOp>(qbarrier, qbarrier.getType(),

  arg_types.push_back(RankedTensorType::get({}, builder.getI32Type()));
  arg_locs.push_back(module_op.getLoc());
  for (const Value input_value : inputs) {
    arg_types.push_back(input_value.getType());
    arg_locs.push_back(input_value.getLoc());
  }

  // Identify result types.
  SmallVector<Type> result_types;
  for (const Value output_value : outputs) {

          return failure();
        }
        scast_op = rewriter.create<quantfork::StorageCastOp>(
            arg.getLoc(), mlir::cast<TensorType>(new_arg_type), arg);
      } else {
        scast_op = rewriter.create<quantfork::StorageCastOp>(
            arg.getLoc(), arg_type.clone(qtype.getStorageType()), arg);
      }
      args.push_back(scast_op.getResult());
    }

        quantize_op->setOperand(quantize_operand_num, dq_op);
        return false;
      }
    }
    rewriter.setInsertionPointAfter(op);
    auto q = rewriter.create<Q>(op->getLoc(), cast_type, op.getResult());
    auto dq = rewriter.create<DQ>(op->getLoc(), expressed_type, q);
    quantize_op->setOperand(quantize_operand_num, dq.getResult());
    return true;
  }

      other.addEntryBlock();
      OpBuilder builder(other.getRegion());
      auto call_op = builder.create<mlir::func::CallOp>(
          f.getLoc(), f.getFunctionType().getResults(), f.getSymName(),
          other.getRegion().getArguments());
      builder.create<mlir::func::ReturnOp>(f.getLoc(), call_op.getResults());
    }

    Unexport(f);
  }
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>>