}
    }
  };

  for (int i = 0, e = func.getNumArguments(); i != e; ++i) {
    BlockArgument arg = func.getArgument(i);
    auto* arg_block = arg.getOwner();
    add_quantize_op(arg.getLoc(), arg.getType(), arg_block,
                    std::next(arg_block->begin(), i), arg, i);
  }

  return false;
}

bool PrepareQuantizePass::RemoveRedundantStats(func::FuncOp func) {

    auto decomposition = mlir::cast<FlatSymbolRefAttr>(calledComputations[0]);

    auto composite = rewriter.create<mlir::stablehlo::CompositeOp>(
        op.getLoc(), op.getResultTypes(), op.getOperands(), name.str(), attrs,
        decomposition.getValue());
    rewriter.replaceOp(op, composite.getResults());
    return success();
  }
};

      if (operand.getType() != expected_type) {
        operand = rewriter.create<TF::CastOp>(
            op.getLoc(), expected_type, operand,
            /*Truncate=*/rewriter.getBoolAttr(false));
      }
      casted_operands.push_back(operand);
    }

    auto call = rewriter.create<func::CallOp>(
        op->getLoc(), main_fn.getSymName(), main_fn.getResultTypes(),
        casted_operands);

    rewriter.setInsertionPointAfter(weight_op);
    auto new_shape_const = rewriter.create<arith::ConstantOp>(
        weight_op->getLoc(), shape_spec_type, new_shape_const_attr);
    auto reshape_op = rewriter.create<TF::ReshapeOp>(
        weight_op->getLoc(), new_shape, weight_op->getResult(0),
        new_shape_const);
    op->setOperand(weight_operand_idx, reshape_op);

      if (!isa<TF::CaseOp>(use)) return;
    }
    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.

    // The "wraps" syntax only encodes a single location.
    // In order to correctly round-trip, we can only use this syntax when all
    // the locations are identical.
    if (wrapped_op.getLoc() == getLoc() && yield_op.getLoc() == getLoc()) {
      p << " wraps ";
      p.printGenericOp(&wrapped_op);
      return;
    }
  }
  p << ' ';
  p.printRegion(getOperation()->getRegion(0));

    if (!input_op) {
      return failure();
    }

    Builder builder(op.getContext());
    auto new_gather_op = rewriter.create<TFL::GatherOp>(
        op.getLoc(),
        /*result=*/
        mlir::cast<TensorType>(op.getResult().getType())
            .clone(builder.getI4Type()),
        /*operand=*/op.getOperands(), op->getAttrs());

  }

  OpBuilder builder(replicate_op);
  // Clone this old replicate op but with new result types.
  auto new_replicate_op = builder.create<tf_device::ReplicateOp>(
      replicate_op->getLoc(), new_result_types, replicate_op->getOperands(),
      replicate_op->getAttrs());

  // Move region to the new op.
  new_replicate_op.getRegion().takeBody(replicate_op.getRegion());

        (target_op_names.empty() && !op->isRegistered())) {
      custom_ops.push_back(op);
    }
  });

  for (auto *op : custom_ops) {
    builder.setInsertionPoint(op);
    Location loc = op->getLoc();
    auto custom_op = builder.create<CustomTfOp>(loc, op->getResultTypes(),
                                                op->getOperands());
    Region region;
    Block *new_block = new Block;

    return failure();
  }

  results.clear();
  builder.setInsertionPointAfter(op);
  for (const auto& result_value : result_values) {
    results.push_back(builder.create<TF::ConstOp>(op->getLoc(), result_value));
  }
  return success();
}

bool IsOperationFoldable(Operation* op) {
  if (isa<TF::ConstOp>(op)) return true;

  if (op->getDialect()->getNamespace() != "tf" || !TF::CanBeFolded(op)) {