Method randn

models/diffusion/base.py:105–121

(
            self,
            *shape,
            rng: List[torch.Generator] | torch.Generator | None = None,
            **kwargs,
    )

        return self._dummy.device

    def randn(
            self,
            *shape,
            rng: List[torch.Generator] | torch.Generator | None = None,
            **kwargs,
    ) -> torch.Tensor:
        if rng is None:
            return torch.randn(*shape, **kwargs)
        elif isinstance(rng, torch.Generator):
            return torch.randn(*shape, generator=rng, **kwargs)
        elif isinstance(rng, list):
            assert len(rng) == shape[0]
            return torch.stack(
                [torch.randn(*shape[1:], generator=r, **kwargs) for r in rng]
            )
        else:
            raise ValueError(f"invalid rng: {rng}")

    def randn_like(
            self,