def METHOD_NAME(self, x):
    B, C, H, W = x.shape
    # (B, C, H, W) -> (B * S, D, N)
    if self.spatial:
        D = C // self.S
        N = H * W
        x = x.reshape([B * self.S, D, N])
    else:
        D = H * W
        N = C // self.S
        x = x.reshape([B * self.S, N, D]).transpose([0, 2, 1])
    if not self.rand_init and not hasattr(self, "bases"):
        bases = self._build_bases(1, self.S, D, self.R)
        self.register_buffer("bases", bases)
    # (S, D, R) -> (B * S, D, R)
    if self.rand_init:
        bases = self._build_bases(B, self.S, D, self.R)
    else:
        bases = paddle.repeat_interleave(self.bases, B, 0)
    bases, coef = self.local_inference(x, bases)
    # (B * S, N, R)
    coef = self.compute_coef(x, bases, coef)
    # (B * S, D, R) @ (B * S, N, R)^T -> (B * S, D, N)
    x = paddle.bmm(bases, coef.transpose([0, 2, 1]))
    # (B * S, D, N) -> (B, C, H, W)
    if self.spatial:
        x = x.reshape([B, C, H, W])
    else:
        x = x.transpose([0, 2, 1]).reshape([B, C, H, W])
    # (B * H, D, R) -> (B, H, N, D)
    # bases = bases.reshape([B, self.S, D, self.R])
    return x