def __call__()

in jcm/models/layers.py [0:0]

• 16 lines of code
• 4 McCabe index (conditional complexity)

    def __call__(self, xs):
        sums = jnp.zeros((xs[0].shape[0], *self.shape, self.features))
        for i in range(len(xs)):
            h = ncsn_conv3x3(xs[i], self.features, stride=1, bias=True)
            if self.interpolation == "bilinear":
                h = jax.image.resize(
                    h, (h.shape[0], *self.shape, h.shape[-1]), "bilinear"
                )
            elif self.interpolation == "nearest_neighbor":
                h = jax.image.resize(
                    h, (h.shape[0], *self.shape, h.shape[-1]), "nearest"
                )
            else:
                raise ValueError(f"Interpolation {self.interpolation} does not exist!")
            sums = sums + h
        return sums