Repository: opencv/opencv
File: modules/dnn/src/layers/affinegrid_layer.cpp (C++)

Code before:
public:
    bool align_corners;

    AffineGridLayerImpl(const LayerParams& params)
    {
        setParamsFrom(params);
        align_corners = params.get<bool>("align_corners", false);
    }

    bool getMemoryShapes(const std::vector<MatShape> &inputs,
                         const int requiredOutputs,
                         std::vector<MatShape> &outputs,
                         std::vector<MatShape> &internals) const CV_OVERRIDE
    {
        CV_CheckGE((int)inputs.size(), 1, "AffineGrid requires at least transform input");

        const MatShape thetaShape = inputs[0];
        CV_CheckGE((int)thetaShape.size(), 3, "theta must be at least 3D [N, r, c]");
        int N = thetaShape[0];
        int r = thetaShape[1];
        int c = thetaShape[2];
        CV_CheckTrue((r == 2 && c == 3) || (r == 3 && c == 4), "theta must be [N,2,3] or [N,3,4]");
        const bool is3d = (r == 3);

        MatShape outShape;
        if ((int)inputs.size() >= 2) {
            const MatShape sizeShape = inputs[1];
            CV_CheckTrue(sizeShape.size() == 1, "size input must be 1D tensor");
            const int dims = (int)sizeShape[0];
            CV_CheckTrue(dims == (is3d ? 5 : 4), "size input must have 4 (2D) or 5 (3D) elements");
            MatShape resolved = is3d ? MatShape{N, -1, -1, -1, 3} : MatShape{N, -1, -1, 2};
            Net::Impl* netimpl_ = getNetImpl(this);
            if (netimpl_) {
                try {
                    Mat sz = netimpl_->argTensor(this->inputs[1]);
                    if (!sz.empty()) {
                        if (!is3d) {
                            int H = -1, W = -1;
                            if (sz.depth() == CV_64S) { H = (int)sz.at<int64_t>(2); W = (int)sz.at<int64_t>(3); }
                            else if (sz.depth() == CV_32S) { H = sz.at<int32_t>(2); W = sz.at<int32_t>(3); }
                            if (H > 0 && W > 0) resolved = MatShape{N, H, W, 2};
                        } else {
                            int D = -1, H = -1, W = -1;
                            if (sz.depth() == CV_64S) { D = (int)sz.at<int64_t>(2); H = (int)sz.at<int64_t>(3); W = (int)sz.at<int64_t>(4); }
                            else if (sz.depth() == CV_32S) { D = sz.at<int32_t>(2); H = sz.at<int32_t>(3); W = sz.at<int32_t>(4); }
                            if (D > 0 && H > 0 && W > 0) resolved = MatShape{N, D, H, W, 3};
                        }
                    }
                } catch (const cv::Exception& e) {
                    CV_Error(cv::Error::StsError, cv::format("DNN/AffineGrid: failed to resolve output shape from 'size' tensor: %s", e.what()));
                }

Review comment (refactor):
I suggest to move this logic into a separate method that would be called from getMemoryShapes if the second input is constant tensor.

Code after:
template<typename T>
void computeGrid2D(const Mat& theta, Mat& grid, int N, int H, int W,
                   float xs, float xd, float ys, float yd)
{
    const int total = N * H;
    parallel_for_(Range(0, total), [&](const Range& range){
        const int H_ = H;
        const int W_ = W;
        const float xs_ = xs, xd_ = xd, ys_ = ys, yd_ = yd;
        const Mat theta_ = theta;
        Mat grid_ = grid;
        int lastN = -1;
        const T* T0 = nullptr;
        const T* T1 = nullptr;
        float T0_0 = 0.f, T1_0 = 0.f;
        for (int nyi = range.start; nyi < range.end; nyi++)
        {
            int n = nyi / H_;
            int y = nyi - n * H_;
            if (n != lastN)
            {
                int i0[3] = {n, 0, 0};
                int i1[3] = {n, 1, 0};
                T0 = theta_.ptr<T>(i0);
                T1 = theta_.ptr<T>(i1);
                T0_0 = (float)T0[0];
                T1_0 = (float)T1[0];
                lastN = n;
            }
            float ny = yd_ + ys_ * y;
            float base0 = (float)T0[1]*ny + (float)T0[2];
            float base1 = (float)T1[1]*ny + (float)T1[2];
            float* out = grid_.ptr<float>(n, y, 0);
            for (int x = 0; x < W_; x++)
            {
                float nx = xd_ + xs_ * x;
                out[2*x + 0] = T0_0*nx + base0;
                out[2*x + 1] = T1_0*nx + base1;
            }
        }
    });
}

template<typename T>
void computeGrid3D(const Mat& theta, Mat& grid, int N, int D, int H, int W,
                   float xs, float xd, float ys, float yd, float zs, float zd)
{
    const int stride = D * H;
    const int total = N * stride;
    parallel_for_(Range(0, total), [&](const Range& range){
        const int H_ = H;