densities.py [161:180]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        locs = [self.manifold.projx(loc) for loc in locs]
        scale = jnp.full(self.manifold.D-1, .3)
        self.dists = [
            WrappedNormal(manifold=self.manifold, loc=loc, scale=scale)
            for loc in locs
        ]

    def log_prob(self, z):
        raise NotImplementedError()

    def sample(self, key, n_samples):
        keys = random.split(key, len(self.dists))
        n = int(np.ceil(n_samples/len(self.dists)))
        samples = jnp.concatenate([
            d.sample(key, n) for key, d in zip(keys, self.dists)
        ], axis=0)
        samples = random.permutation(key, samples)
        return samples[:n_samples]

    def __hash__(self): return 0 # For jitting
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



densities.py [190:209]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        locs = [self.manifold.projx(loc) for loc in locs]
        scale = jnp.full(self.manifold.D-1, .3)
        self.dists = [
            WrappedNormal(manifold=self.manifold, loc=loc, scale=scale)
            for loc in locs
        ]

    def log_prob(self, z):
        raise NotImplementedError()

    def sample(self, key, n_samples):
        keys = random.split(key, len(self.dists))
        n = int(np.ceil(n_samples/len(self.dists)))
        samples = jnp.concatenate([
            d.sample(key, n) for key, d in zip(keys, self.dists)
        ], axis=0)
        samples = random.permutation(key, samples)
        return samples[:n_samples]

    def __hash__(self): return 0 # For jitting
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -