diff --git "a/functions_data.csv" "b/functions_data.csv"
new file mode 100644--- /dev/null
+++ "b/functions_data.csv"
@@ -0,0 +1,3369 @@
+function_name,docstring,function_body,file_path
+diffusion_from_config,,"def diffusion_from_config(config: Dict[str, Any]) ->GaussianDiffusion:
+    schedule = config['schedule']
+    steps = config['timesteps']
+    respace = config.get('respacing', None)
+    mean_type = config.get('mean_type', 'epsilon')
+    betas = get_named_beta_schedule(schedule, steps)
+    channel_scales = config.get('channel_scales', None)
+    channel_biases = config.get('channel_biases', None)
+    if channel_scales is not None:
+        channel_scales = np.array(channel_scales)
+    if channel_biases is not None:
+        channel_biases = np.array(channel_biases)
+    kwargs = dict(betas=betas, model_mean_type=mean_type, model_var_type=
+        'learned_range', loss_type='mse', channel_scales=channel_scales,
+        channel_biases=channel_biases)
+    if respace is None:
+        return GaussianDiffusion(**kwargs)
+    else:
+        return SpacedDiffusion(use_timesteps=space_timesteps(steps, respace
+            ), **kwargs)
+",point_e\diffusion\configs.py
+get_beta_schedule,"This is the deprecated API for creating beta schedules.
+
+See get_named_beta_schedule() for the new library of schedules.","def get_beta_schedule(beta_schedule, *, beta_start, beta_end,
+    num_diffusion_timesteps):
+    """"""""""""
+    if beta_schedule == 'linear':
+        betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps,
+            dtype=np.float64)
+    else:
+        raise NotImplementedError(beta_schedule)
+    assert betas.shape == (num_diffusion_timesteps,)
+    return betas
+",point_e\diffusion\gaussian_diffusion.py
+get_named_beta_schedule,"Get a pre-defined beta schedule for the given name.
+
+The beta schedule library consists of beta schedules which remain similar
+in the limit of num_diffusion_timesteps.
+Beta schedules may be added, but should not be removed or changed once
+they are committed to maintain backwards compatibility.","def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
+    """"""""""""
+    if schedule_name == 'linear':
+        scale = 1000 / num_diffusion_timesteps
+        return get_beta_schedule('linear', beta_start=scale * 0.0001,
+            beta_end=scale * 0.02, num_diffusion_timesteps=
+            num_diffusion_timesteps)
+    elif schedule_name == 'cosine':
+        return betas_for_alpha_bar(num_diffusion_timesteps, lambda t: math.
+            cos((t + 0.008) / 1.008 * math.pi / 2) ** 2)
+    else:
+        raise NotImplementedError(f'unknown beta schedule: {schedule_name}')
+",point_e\diffusion\gaussian_diffusion.py
+betas_for_alpha_bar,"Create a beta schedule that discretizes the given alpha_t_bar function,
+which defines the cumulative product of (1-beta) over time from t = [0,1].
+
+:param num_diffusion_timesteps: the number of betas to produce.
+:param alpha_bar: a lambda that takes an argument t from 0 to 1 and
+                  produces the cumulative product of (1-beta) up to that
+                  part of the diffusion process.
+:param max_beta: the maximum beta to use; use values lower than 1 to
+                 prevent singularities.","def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
+    """"""""""""
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+    return np.array(betas)
+",point_e\diffusion\gaussian_diffusion.py
+space_timesteps,"Create a list of timesteps to use from an original diffusion process,
+given the number of timesteps we want to take from equally-sized portions
+of the original process.
+For example, if there's 300 timesteps and the section counts are [10,15,20]
+then the first 100 timesteps are strided to be 10 timesteps, the second 100
+are strided to be 15 timesteps, and the final 100 are strided to be 20.
+:param num_timesteps: the number of diffusion steps in the original
+                      process to divide up.
+:param section_counts: either a list of numbers, or a string containing
+                       comma-separated numbers, indicating the step count
+                       per section. As a special case, use ""ddimN"" where N
+                       is a number of steps to use the striding from the
+                       DDIM paper.
+:return: a set of diffusion steps from the original process to use.","def space_timesteps(num_timesteps, section_counts):
+    """"""""""""
+    if isinstance(section_counts, str):
+        if section_counts.startswith('ddim'):
+            desired_count = int(section_counts[len('ddim'):])
+            for i in range(1, num_timesteps):
+                if len(range(0, num_timesteps, i)) == desired_count:
+                    return set(range(0, num_timesteps, i))
+            raise ValueError(
+                f'cannot create exactly {num_timesteps} steps with an integer stride'
+                )
+        elif section_counts.startswith('exact'):
+            res = set(int(x) for x in section_counts[len('exact'):].split(','))
+            for x in res:
+                if x < 0 or x >= num_timesteps:
+                    raise ValueError(f'timestep out of bounds: {x}')
+            return res
+        section_counts = [int(x) for x in section_counts.split(',')]
+    size_per = num_timesteps // len(section_counts)
+    extra = num_timesteps % len(section_counts)
+    start_idx = 0
+    all_steps = []
+    for i, section_count in enumerate(section_counts):
+        size = size_per + (1 if i < extra else 0)
+        if size < section_count:
+            raise ValueError(
+                f'cannot divide section of {size} steps into {section_count}')
+        if section_count <= 1:
+            frac_stride = 1
+        else:
+            frac_stride = (size - 1) / (section_count - 1)
+        cur_idx = 0.0
+        taken_steps = []
+        for _ in range(section_count):
+            taken_steps.append(start_idx + round(cur_idx))
+            cur_idx += frac_stride
+        all_steps += taken_steps
+        start_idx += size
+    return set(all_steps)
+",point_e\diffusion\gaussian_diffusion.py
+_extract_into_tensor,"Extract values from a 1-D numpy array for a batch of indices.
+
+:param arr: the 1-D numpy array.
+:param timesteps: a tensor of indices into the array to extract.
+:param broadcast_shape: a larger shape of K dimensions with the batch
+                        dimension equal to the length of timesteps.
+:return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.","def _extract_into_tensor(arr, timesteps, broadcast_shape):
+    """"""""""""
+    res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
+    while len(res.shape) < len(broadcast_shape):
+        res = res[..., None]
+    return res + th.zeros(broadcast_shape, device=timesteps.device)
+",point_e\diffusion\gaussian_diffusion.py
+normal_kl,"Compute the KL divergence between two gaussians.
+Shapes are automatically broadcasted, so batches can be compared to
+scalars, among other use cases.","def normal_kl(mean1, logvar1, mean2, logvar2):
+    """"""""""""
+    tensor = None
+    for obj in (mean1, logvar1, mean2, logvar2):
+        if isinstance(obj, th.Tensor):
+            tensor = obj
+            break
+    assert tensor is not None, 'at least one argument must be a Tensor'
+    logvar1, logvar2 = [(x if isinstance(x, th.Tensor) else th.tensor(x).to
+        (tensor)) for x in (logvar1, logvar2)]
+    return 0.5 * (-1.0 + logvar2 - logvar1 + th.exp(logvar1 - logvar2) + (
+        mean1 - mean2) ** 2 * th.exp(-logvar2))
+",point_e\diffusion\gaussian_diffusion.py
+approx_standard_normal_cdf,"A fast approximation of the cumulative distribution function of the
+standard normal.","def approx_standard_normal_cdf(x):
+    """"""""""""
+    return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.
+        pow(x, 3))))
+",point_e\diffusion\gaussian_diffusion.py
+discretized_gaussian_log_likelihood,"Compute the log-likelihood of a Gaussian distribution discretizing to a
+given image.
+:param x: the target images. It is assumed that this was uint8 values,
+          rescaled to the range [-1, 1].
+:param means: the Gaussian mean Tensor.
+:param log_scales: the Gaussian log stddev Tensor.
+:return: a tensor like x of log probabilities (in nats).","def discretized_gaussian_log_likelihood(x, *, means, log_scales):
+    """"""""""""
+    assert x.shape == means.shape == log_scales.shape
+    centered_x = x - means
+    inv_stdv = th.exp(-log_scales)
+    plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
+    cdf_plus = approx_standard_normal_cdf(plus_in)
+    min_in = inv_stdv * (centered_x - 1.0 / 255.0)
+    cdf_min = approx_standard_normal_cdf(min_in)
+    log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))
+    log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))
+    cdf_delta = cdf_plus - cdf_min
+    log_probs = th.where(x < -0.999, log_cdf_plus, th.where(x > 0.999,
+        log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))))
+    assert log_probs.shape == x.shape
+    return log_probs
+",point_e\diffusion\gaussian_diffusion.py
+mean_flat,Take the mean over all non-batch dimensions.,"def mean_flat(tensor):
+    """"""""""""
+    return tensor.flatten(1).mean(1)
+",point_e\diffusion\gaussian_diffusion.py
+__init__,,"def __init__(self, *, betas: Sequence[float], model_mean_type: str,
+    model_var_type: str, loss_type: str, discretized_t0: bool=False,
+    channel_scales: Optional[np.ndarray]=None, channel_biases: Optional[np.
+    ndarray]=None):
+    self.model_mean_type = model_mean_type
+    self.model_var_type = model_var_type
+    self.loss_type = loss_type
+    self.discretized_t0 = discretized_t0
+    self.channel_scales = channel_scales
+    self.channel_biases = channel_biases
+    betas = np.array(betas, dtype=np.float64)
+    self.betas = betas
+    assert len(betas.shape) == 1, 'betas must be 1-D'
+    assert (betas > 0).all() and (betas <= 1).all()
+    self.num_timesteps = int(betas.shape[0])
+    alphas = 1.0 - betas
+    self.alphas_cumprod = np.cumprod(alphas, axis=0)
+    self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
+    self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
+    assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)
+    self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
+    self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
+    self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
+    self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
+    self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
+    self.posterior_variance = betas * (1.0 - self.alphas_cumprod_prev) / (
+        1.0 - self.alphas_cumprod)
+    self.posterior_log_variance_clipped = np.log(np.append(self.
+        posterior_variance[1], self.posterior_variance[1:]))
+    self.posterior_mean_coef1 = betas * np.sqrt(self.alphas_cumprod_prev) / (
+        1.0 - self.alphas_cumprod)
+    self.posterior_mean_coef2 = (1.0 - self.alphas_cumprod_prev) * np.sqrt(
+        alphas) / (1.0 - self.alphas_cumprod)
+",point_e\diffusion\gaussian_diffusion.py
+get_sigmas,,"def get_sigmas(self, t):
+    return _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, t.shape)
+",point_e\diffusion\gaussian_diffusion.py
+q_mean_variance,"Get the distribution q(x_t | x_0).
+
+:param x_start: the [N x C x ...] tensor of noiseless inputs.
+:param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+:return: A tuple (mean, variance, log_variance), all of x_start's shape.","def q_mean_variance(self, x_start, t):
+    """"""""""""
+    mean = _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape
+        ) * x_start
+    variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape
+        )
+    log_variance = _extract_into_tensor(self.log_one_minus_alphas_cumprod,
+        t, x_start.shape)
+    return mean, variance, log_variance
+",point_e\diffusion\gaussian_diffusion.py
+q_sample,"Diffuse the data for a given number of diffusion steps.
+
+In other words, sample from q(x_t | x_0).
+
+:param x_start: the initial data batch.
+:param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+:param noise: if specified, the split-out normal noise.
+:return: A noisy version of x_start.","def q_sample(self, x_start, t, noise=None):
+    """"""""""""
+    if noise is None:
+        noise = th.randn_like(x_start)
+    assert noise.shape == x_start.shape
+    return _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape
+        ) * x_start + _extract_into_tensor(self.
+        sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
+",point_e\diffusion\gaussian_diffusion.py
+q_posterior_mean_variance,"Compute the mean and variance of the diffusion posterior:
+
+    q(x_{t-1} | x_t, x_0)","def q_posterior_mean_variance(self, x_start, x_t, t):
+    """"""""""""
+    assert x_start.shape == x_t.shape
+    posterior_mean = _extract_into_tensor(self.posterior_mean_coef1, t, x_t
+        .shape) * x_start + _extract_into_tensor(self.posterior_mean_coef2,
+        t, x_t.shape) * x_t
+    posterior_variance = _extract_into_tensor(self.posterior_variance, t,
+        x_t.shape)
+    posterior_log_variance_clipped = _extract_into_tensor(self.
+        posterior_log_variance_clipped, t, x_t.shape)
+    assert posterior_mean.shape[0] == posterior_variance.shape[0
+        ] == posterior_log_variance_clipped.shape[0] == x_start.shape[0]
+    return posterior_mean, posterior_variance, posterior_log_variance_clipped
+",point_e\diffusion\gaussian_diffusion.py
+p_mean_variance,"Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
+the initial x, x_0.
+
+:param model: the model, which takes a signal and a batch of timesteps
+              as input.
+:param x: the [N x C x ...] tensor at time t.
+:param t: a 1-D Tensor of timesteps.
+:param clip_denoised: if True, clip the denoised signal into [-1, 1].
+:param denoised_fn: if not None, a function which applies to the
+    x_start prediction before it is used to sample. Applies before
+    clip_denoised.
+:param model_kwargs: if not None, a dict of extra keyword arguments to
+    pass to the model. This can be used for conditioning.
+:return: a dict with the following keys:
+         - 'mean': the model mean output.
+         - 'variance': the model variance output.
+         - 'log_variance': the log of 'variance'.
+         - 'pred_xstart': the prediction for x_0.","def p_mean_variance(self, model, x, t, clip_denoised=False, denoised_fn=
+    None, model_kwargs=None):
+    """"""""""""
+    if model_kwargs is None:
+        model_kwargs = {}
+    B, C = x.shape[:2]
+    assert t.shape == (B,)
+    model_output = model(x, t, **model_kwargs)
+    if isinstance(model_output, tuple):
+        model_output, extra = model_output
+    else:
+        extra = None
+    if self.model_var_type in ['learned', 'learned_range']:
+        assert model_output.shape == (B, C * 2, *x.shape[2:])
+        model_output, model_var_values = th.split(model_output, C, dim=1)
+        if self.model_var_type == 'learned':
+            model_log_variance = model_var_values
+            model_variance = th.exp(model_log_variance)
+        else:
+            min_log = _extract_into_tensor(self.
+                posterior_log_variance_clipped, t, x.shape)
+            max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
+            frac = (model_var_values + 1) / 2
+            model_log_variance = frac * max_log + (1 - frac) * min_log
+            model_variance = th.exp(model_log_variance)
+    else:
+        model_variance, model_log_variance = {'fixed_large': (np.append(
+            self.posterior_variance[1], self.betas[1:]), np.log(np.append(
+            self.posterior_variance[1], self.betas[1:]))), 'fixed_small': (
+            self.posterior_variance, self.posterior_log_variance_clipped)}[self
+            .model_var_type]
+        model_variance = _extract_into_tensor(model_variance, t, x.shape)
+        model_log_variance = _extract_into_tensor(model_log_variance, t, x.
+            shape)
+
+    def process_xstart(x):
+        if denoised_fn is not None:
+            x = denoised_fn(x)
+        if clip_denoised:
+            return x.clamp(-1, 1)
+        return x
+    if self.model_mean_type == 'x_prev':
+        pred_xstart = process_xstart(self._predict_xstart_from_xprev(x_t=x,
+            t=t, xprev=model_output))
+        model_mean = model_output
+    elif self.model_mean_type in ['x_start', 'epsilon']:
+        if self.model_mean_type == 'x_start':
+            pred_xstart = process_xstart(model_output)
+        else:
+            pred_xstart = process_xstart(self._predict_xstart_from_eps(x_t=
+                x, t=t, eps=model_output))
+        model_mean, _, _ = self.q_posterior_mean_variance(x_start=
+            pred_xstart, x_t=x, t=t)
+    else:
+        raise NotImplementedError(self.model_mean_type)
+    assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
+    return {'mean': model_mean, 'variance': model_variance, 'log_variance':
+        model_log_variance, 'pred_xstart': pred_xstart, 'extra': extra}
+",point_e\diffusion\gaussian_diffusion.py
+_predict_xstart_from_eps,,"def _predict_xstart_from_eps(self, x_t, t, eps):
+    assert x_t.shape == eps.shape
+    return _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape
+        ) * x_t - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t,
+        x_t.shape) * eps
+",point_e\diffusion\gaussian_diffusion.py
+_predict_xstart_from_xprev,,"def _predict_xstart_from_xprev(self, x_t, t, xprev):
+    assert x_t.shape == xprev.shape
+    return _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape
+        ) * xprev - _extract_into_tensor(self.posterior_mean_coef2 / self.
+        posterior_mean_coef1, t, x_t.shape) * x_t
+",point_e\diffusion\gaussian_diffusion.py
+_predict_eps_from_xstart,,"def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
+    return (_extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.
+        shape) * x_t - pred_xstart) / _extract_into_tensor(self.
+        sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+",point_e\diffusion\gaussian_diffusion.py
+condition_mean,"Compute the mean for the previous step, given a function cond_fn that
+computes the gradient of a conditional log probability with respect to
+x. In particular, cond_fn computes grad(log(p(y|x))), and we want to
+condition on y.
+
+This uses the conditioning strategy from Sohl-Dickstein et al. (2015).","def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+    """"""""""""
+    gradient = cond_fn(x, t, **model_kwargs)
+    new_mean = p_mean_var['mean'].float() + p_mean_var['variance'
+        ] * gradient.float()
+    return new_mean
+",point_e\diffusion\gaussian_diffusion.py
+condition_score,"Compute what the p_mean_variance output would have been, should the
+model's score function be conditioned by cond_fn.
+
+See condition_mean() for details on cond_fn.
+
+Unlike condition_mean(), this instead uses the conditioning strategy
+from Song et al (2020).","def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+    """"""""""""
+    alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+    eps = self._predict_eps_from_xstart(x, t, p_mean_var['pred_xstart'])
+    eps = eps - (1 - alpha_bar).sqrt() * cond_fn(x, t, **model_kwargs)
+    out = p_mean_var.copy()
+    out['pred_xstart'] = self._predict_xstart_from_eps(x, t, eps)
+    out['mean'], _, _ = self.q_posterior_mean_variance(x_start=out[
+        'pred_xstart'], x_t=x, t=t)
+    return out
+",point_e\diffusion\gaussian_diffusion.py
+p_sample,"Sample x_{t-1} from the model at the given timestep.
+
+:param model: the model to sample from.
+:param x: the current tensor at x_{t-1}.
+:param t: the value of t, starting at 0 for the first diffusion step.
+:param clip_denoised: if True, clip the x_start prediction to [-1, 1].
+:param denoised_fn: if not None, a function which applies to the
+    x_start prediction before it is used to sample.
+:param cond_fn: if not None, this is a gradient function that acts
+                similarly to the model.
+:param model_kwargs: if not None, a dict of extra keyword arguments to
+    pass to the model. This can be used for conditioning.
+:return: a dict containing the following keys:
+         - 'sample': a random sample from the model.
+         - 'pred_xstart': a prediction of x_0.","def p_sample(self, model, x, t, clip_denoised=False, denoised_fn=None,
+    cond_fn=None, model_kwargs=None):
+    """"""""""""
+    out = self.p_mean_variance(model, x, t, clip_denoised=clip_denoised,
+        denoised_fn=denoised_fn, model_kwargs=model_kwargs)
+    noise = th.randn_like(x)
+    nonzero_mask = (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+    if cond_fn is not None:
+        out['mean'] = self.condition_mean(cond_fn, out, x, t, model_kwargs=
+            model_kwargs)
+    sample = out['mean'] + nonzero_mask * th.exp(0.5 * out['log_variance']
+        ) * noise
+    return {'sample': sample, 'pred_xstart': out['pred_xstart']}
+",point_e\diffusion\gaussian_diffusion.py
+p_sample_loop,"Generate samples from the model.
+
+:param model: the model module.
+:param shape: the shape of the samples, (N, C, H, W).
+:param noise: if specified, the noise from the encoder to sample.
+              Should be of the same shape as `shape`.
+:param clip_denoised: if True, clip x_start predictions to [-1, 1].
+:param denoised_fn: if not None, a function which applies to the
+    x_start prediction before it is used to sample.
+:param cond_fn: if not None, this is a gradient function that acts
+                similarly to the model.
+:param model_kwargs: if not None, a dict of extra keyword arguments to
+    pass to the model. This can be used for conditioning.
+:param device: if specified, the device to create the samples on.
+               If not specified, use a model parameter's device.
+:param progress: if True, show a tqdm progress bar.
+:return: a non-differentiable batch of samples.","def p_sample_loop(self, model, shape, noise=None, clip_denoised=False,
+    denoised_fn=None, cond_fn=None, model_kwargs=None, device=None,
+    progress=False, temp=1.0):
+    """"""""""""
+    final = None
+    for sample in self.p_sample_loop_progressive(model, shape, noise=noise,
+        clip_denoised=clip_denoised, denoised_fn=denoised_fn, cond_fn=
+        cond_fn, model_kwargs=model_kwargs, device=device, progress=
+        progress, temp=temp):
+        final = sample
+    return final['sample']
+",point_e\diffusion\gaussian_diffusion.py
+p_sample_loop_progressive,"Generate samples from the model and yield intermediate samples from
+each timestep of diffusion.
+
+Arguments are the same as p_sample_loop().
+Returns a generator over dicts, where each dict is the return value of
+p_sample().","def p_sample_loop_progressive(self, model, shape, noise=None, clip_denoised
+    =False, denoised_fn=None, cond_fn=None, model_kwargs=None, device=None,
+    progress=False, temp=1.0):
+    """"""""""""
+    if device is None:
+        device = next(model.parameters()).device
+    assert isinstance(shape, (tuple, list))
+    if noise is not None:
+        img = noise
+    else:
+        img = th.randn(*shape, device=device) * temp
+    indices = list(range(self.num_timesteps))[::-1]
+    if progress:
+        from tqdm.auto import tqdm
+        indices = tqdm(indices)
+    for i in indices:
+        t = th.tensor([i] * shape[0], device=device)
+        with th.no_grad():
+            out = self.p_sample(model, img, t, clip_denoised=clip_denoised,
+                denoised_fn=denoised_fn, cond_fn=cond_fn, model_kwargs=
+                model_kwargs)
+            yield self.unscale_out_dict(out)
+            img = out['sample']
+",point_e\diffusion\gaussian_diffusion.py
+ddim_sample,"Sample x_{t-1} from the model using DDIM.
+
+Same usage as p_sample().","def ddim_sample(self, model, x, t, clip_denoised=False, denoised_fn=None,
+    cond_fn=None, model_kwargs=None, eta=0.0):
+    """"""""""""
+    out = self.p_mean_variance(model, x, t, clip_denoised=clip_denoised,
+        denoised_fn=denoised_fn, model_kwargs=model_kwargs)
+    if cond_fn is not None:
+        out = self.condition_score(cond_fn, out, x, t, model_kwargs=
+            model_kwargs)
+    eps = self._predict_eps_from_xstart(x, t, out['pred_xstart'])
+    alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+    alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
+    sigma = eta * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) * th.sqrt(
+        1 - alpha_bar / alpha_bar_prev)
+    noise = th.randn_like(x)
+    mean_pred = out['pred_xstart'] * th.sqrt(alpha_bar_prev) + th.sqrt(1 -
+        alpha_bar_prev - sigma ** 2) * eps
+    nonzero_mask = (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+    sample = mean_pred + nonzero_mask * sigma * noise
+    return {'sample': sample, 'pred_xstart': out['pred_xstart']}
+",point_e\diffusion\gaussian_diffusion.py
+ddim_reverse_sample,Sample x_{t+1} from the model using DDIM reverse ODE.,"def ddim_reverse_sample(self, model, x, t, clip_denoised=False, denoised_fn
+    =None, cond_fn=None, model_kwargs=None, eta=0.0):
+    """"""""""""
+    assert eta == 0.0, 'Reverse ODE only for deterministic path'
+    out = self.p_mean_variance(model, x, t, clip_denoised=clip_denoised,
+        denoised_fn=denoised_fn, model_kwargs=model_kwargs)
+    if cond_fn is not None:
+        out = self.condition_score(cond_fn, out, x, t, model_kwargs=
+            model_kwargs)
+    eps = (_extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) *
+        x - out['pred_xstart']) / _extract_into_tensor(self.
+        sqrt_recipm1_alphas_cumprod, t, x.shape)
+    alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)
+    mean_pred = out['pred_xstart'] * th.sqrt(alpha_bar_next) + th.sqrt(1 -
+        alpha_bar_next) * eps
+    return {'sample': mean_pred, 'pred_xstart': out['pred_xstart']}
+",point_e\diffusion\gaussian_diffusion.py
+ddim_sample_loop,"Generate samples from the model using DDIM.
+
+Same usage as p_sample_loop().","def ddim_sample_loop(self, model, shape, noise=None, clip_denoised=False,
+    denoised_fn=None, cond_fn=None, model_kwargs=None, device=None,
+    progress=False, eta=0.0, temp=1.0):
+    """"""""""""
+    final = None
+    for sample in self.ddim_sample_loop_progressive(model, shape, noise=
+        noise, clip_denoised=clip_denoised, denoised_fn=denoised_fn,
+        cond_fn=cond_fn, model_kwargs=model_kwargs, device=device, progress
+        =progress, eta=eta, temp=temp):
+        final = sample
+    return final['sample']
+",point_e\diffusion\gaussian_diffusion.py
+ddim_sample_loop_progressive,"Use DDIM to sample from the model and yield intermediate samples from
+each timestep of DDIM.
+
+Same usage as p_sample_loop_progressive().","def ddim_sample_loop_progressive(self, model, shape, noise=None,
+    clip_denoised=False, denoised_fn=None, cond_fn=None, model_kwargs=None,
+    device=None, progress=False, eta=0.0, temp=1.0):
+    """"""""""""
+    if device is None:
+        device = next(model.parameters()).device
+    assert isinstance(shape, (tuple, list))
+    if noise is not None:
+        img = noise
+    else:
+        img = th.randn(*shape, device=device) * temp
+    indices = list(range(self.num_timesteps))[::-1]
+    if progress:
+        from tqdm.auto import tqdm
+        indices = tqdm(indices)
+    for i in indices:
+        t = th.tensor([i] * shape[0], device=device)
+        with th.no_grad():
+            out = self.ddim_sample(model, img, t, clip_denoised=
+                clip_denoised, denoised_fn=denoised_fn, cond_fn=cond_fn,
+                model_kwargs=model_kwargs, eta=eta)
+            yield self.unscale_out_dict(out)
+            img = out['sample']
+",point_e\diffusion\gaussian_diffusion.py
+_vb_terms_bpd,"Get a term for the variational lower-bound.
+
+The resulting units are bits (rather than nats, as one might expect).
+This allows for comparison to other papers.
+
+:return: a dict with the following keys:
+         - 'output': a shape [N] tensor of NLLs or KLs.
+         - 'pred_xstart': the x_0 predictions.","def _vb_terms_bpd(self, model, x_start, x_t, t, clip_denoised=False,
+    model_kwargs=None):
+    """"""""""""
+    true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
+        x_start=x_start, x_t=x_t, t=t)
+    out = self.p_mean_variance(model, x_t, t, clip_denoised=clip_denoised,
+        model_kwargs=model_kwargs)
+    kl = normal_kl(true_mean, true_log_variance_clipped, out['mean'], out[
+        'log_variance'])
+    kl = mean_flat(kl) / np.log(2.0)
+    decoder_nll = -discretized_gaussian_log_likelihood(x_start, means=out[
+        'mean'], log_scales=0.5 * out['log_variance'])
+    if not self.discretized_t0:
+        decoder_nll = th.zeros_like(decoder_nll)
+    assert decoder_nll.shape == x_start.shape
+    decoder_nll = mean_flat(decoder_nll) / np.log(2.0)
+    output = th.where(t == 0, decoder_nll, kl)
+    return {'output': output, 'pred_xstart': out['pred_xstart'], 'extra':
+        out['extra']}
+",point_e\diffusion\gaussian_diffusion.py
+training_losses,"Compute training losses for a single timestep.
+
+:param model: the model to evaluate loss on.
+:param x_start: the [N x C x ...] tensor of inputs.
+:param t: a batch of timestep indices.
+:param model_kwargs: if not None, a dict of extra keyword arguments to
+    pass to the model. This can be used for conditioning.
+:param noise: if specified, the specific Gaussian noise to try to remove.
+:return: a dict with the key ""loss"" containing a tensor of shape [N].
+         Some mean or variance settings may also have other keys.","def training_losses(self, model, x_start, t, model_kwargs=None, noise=None
+    ) ->Dict[str, th.Tensor]:
+    """"""""""""
+    x_start = self.scale_channels(x_start)
+    if model_kwargs is None:
+        model_kwargs = {}
+    if noise is None:
+        noise = th.randn_like(x_start)
+    x_t = self.q_sample(x_start, t, noise=noise)
+    terms = {}
+    if self.loss_type == 'kl' or self.loss_type == 'rescaled_kl':
+        vb_terms = self._vb_terms_bpd(model=model, x_start=x_start, x_t=x_t,
+            t=t, clip_denoised=False, model_kwargs=model_kwargs)
+        terms['loss'] = vb_terms['output']
+        if self.loss_type == 'rescaled_kl':
+            terms['loss'] *= self.num_timesteps
+        extra = vb_terms['extra']
+    elif self.loss_type == 'mse' or self.loss_type == 'rescaled_mse':
+        model_output = model(x_t, t, **model_kwargs)
+        if isinstance(model_output, tuple):
+            model_output, extra = model_output
+        else:
+            extra = {}
+        if self.model_var_type in ['learned', 'learned_range']:
+            B, C = x_t.shape[:2]
+            assert model_output.shape == (B, C * 2, *x_t.shape[2:])
+            model_output, model_var_values = th.split(model_output, C, dim=1)
+            frozen_out = th.cat([model_output.detach(), model_var_values],
+                dim=1)
+            terms['vb'] = self._vb_terms_bpd(model=lambda *args, r=
+                frozen_out: r, x_start=x_start, x_t=x_t, t=t, clip_denoised
+                =False)['output']
+            if self.loss_type == 'rescaled_mse':
+                terms['vb'] *= self.num_timesteps / 1000.0
+        target = {'x_prev': self.q_posterior_mean_variance(x_start=x_start,
+            x_t=x_t, t=t)[0], 'x_start': x_start, 'epsilon': noise}[self.
+            model_mean_type]
+        assert model_output.shape == target.shape == x_start.shape
+        terms['mse'] = mean_flat((target - model_output) ** 2)
+        if 'vb' in terms:
+            terms['loss'] = terms['mse'] + terms['vb']
+        else:
+            terms['loss'] = terms['mse']
+    else:
+        raise NotImplementedError(self.loss_type)
+    if 'losses' in extra:
+        terms.update({k: loss for k, (loss, _scale) in extra['losses'].items()}
+            )
+        for loss, scale in extra['losses'].values():
+            terms['loss'] = terms['loss'] + loss * scale
+    return terms
+",point_e\diffusion\gaussian_diffusion.py
+_prior_bpd,"Get the prior KL term for the variational lower-bound, measured in
+bits-per-dim.
+
+This term can't be optimized, as it only depends on the encoder.
+
+:param x_start: the [N x C x ...] tensor of inputs.
+:return: a batch of [N] KL values (in bits), one per batch element.","def _prior_bpd(self, x_start):
+    """"""""""""
+    batch_size = x_start.shape[0]
+    t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
+    qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
+    kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0,
+        logvar2=0.0)
+    return mean_flat(kl_prior) / np.log(2.0)
+",point_e\diffusion\gaussian_diffusion.py
+calc_bpd_loop,"Compute the entire variational lower-bound, measured in bits-per-dim,
+as well as other related quantities.
+
+:param model: the model to evaluate loss on.
+:param x_start: the [N x C x ...] tensor of inputs.
+:param clip_denoised: if True, clip denoised samples.
+:param model_kwargs: if not None, a dict of extra keyword arguments to
+    pass to the model. This can be used for conditioning.
+
+:return: a dict containing the following keys:
+         - total_bpd: the total variational lower-bound, per batch element.
+         - prior_bpd: the prior term in the lower-bound.
+         - vb: an [N x T] tensor of terms in the lower-bound.
+         - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.
+         - mse: an [N x T] tensor of epsilon MSEs for each timestep.","def calc_bpd_loop(self, model, x_start, clip_denoised=False, model_kwargs=None
+    ):
+    """"""""""""
+    device = x_start.device
+    batch_size = x_start.shape[0]
+    vb = []
+    xstart_mse = []
+    mse = []
+    for t in list(range(self.num_timesteps))[::-1]:
+        t_batch = th.tensor([t] * batch_size, device=device)
+        noise = th.randn_like(x_start)
+        x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)
+        with th.no_grad():
+            out = self._vb_terms_bpd(model, x_start=x_start, x_t=x_t, t=
+                t_batch, clip_denoised=clip_denoised, model_kwargs=model_kwargs
+                )
+        vb.append(out['output'])
+        xstart_mse.append(mean_flat((out['pred_xstart'] - x_start) ** 2))
+        eps = self._predict_eps_from_xstart(x_t, t_batch, out['pred_xstart'])
+        mse.append(mean_flat((eps - noise) ** 2))
+    vb = th.stack(vb, dim=1)
+    xstart_mse = th.stack(xstart_mse, dim=1)
+    mse = th.stack(mse, dim=1)
+    prior_bpd = self._prior_bpd(x_start)
+    total_bpd = vb.sum(dim=1) + prior_bpd
+    return {'total_bpd': total_bpd, 'prior_bpd': prior_bpd, 'vb': vb,
+        'xstart_mse': xstart_mse, 'mse': mse}
+",point_e\diffusion\gaussian_diffusion.py
+scale_channels,,"def scale_channels(self, x: th.Tensor) ->th.Tensor:
+    if self.channel_scales is not None:
+        x = x * th.from_numpy(self.channel_scales).to(x).reshape([1, -1, *(
+            [1] * (len(x.shape) - 2))])
+    if self.channel_biases is not None:
+        x = x + th.from_numpy(self.channel_biases).to(x).reshape([1, -1, *(
+            [1] * (len(x.shape) - 2))])
+    return x
+",point_e\diffusion\gaussian_diffusion.py
+unscale_channels,,"def unscale_channels(self, x: th.Tensor) ->th.Tensor:
+    if self.channel_biases is not None:
+        x = x - th.from_numpy(self.channel_biases).to(x).reshape([1, -1, *(
+            [1] * (len(x.shape) - 2))])
+    if self.channel_scales is not None:
+        x = x / th.from_numpy(self.channel_scales).to(x).reshape([1, -1, *(
+            [1] * (len(x.shape) - 2))])
+    return x
+",point_e\diffusion\gaussian_diffusion.py
+unscale_out_dict,,"def unscale_out_dict(self, out: Dict[str, Union[th.Tensor, Any]]) ->Dict[
+    str, Union[th.Tensor, Any]]:
+    return {k: (self.unscale_channels(v) if isinstance(v, th.Tensor) else v
+        ) for k, v in out.items()}
+",point_e\diffusion\gaussian_diffusion.py
+__init__,,"def __init__(self, use_timesteps: Iterable[int], **kwargs):
+    self.use_timesteps = set(use_timesteps)
+    self.timestep_map = []
+    self.original_num_steps = len(kwargs['betas'])
+    base_diffusion = GaussianDiffusion(**kwargs)
+    last_alpha_cumprod = 1.0
+    new_betas = []
+    for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):
+        if i in self.use_timesteps:
+            new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)
+            last_alpha_cumprod = alpha_cumprod
+            self.timestep_map.append(i)
+    kwargs['betas'] = np.array(new_betas)
+    super().__init__(**kwargs)
+",point_e\diffusion\gaussian_diffusion.py
+p_mean_variance,,"def p_mean_variance(self, model, *args, **kwargs):
+    return super().p_mean_variance(self._wrap_model(model), *args, **kwargs)
+",point_e\diffusion\gaussian_diffusion.py
+training_losses,,"def training_losses(self, model, *args, **kwargs):
+    return super().training_losses(self._wrap_model(model), *args, **kwargs)
+",point_e\diffusion\gaussian_diffusion.py
+condition_mean,,"def condition_mean(self, cond_fn, *args, **kwargs):
+    return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs)
+",point_e\diffusion\gaussian_diffusion.py
+condition_score,,"def condition_score(self, cond_fn, *args, **kwargs):
+    return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs)
+",point_e\diffusion\gaussian_diffusion.py
+_wrap_model,,"def _wrap_model(self, model):
+    if isinstance(model, _WrappedModel):
+        return model
+    return _WrappedModel(model, self.timestep_map, self.original_num_steps)
+",point_e\diffusion\gaussian_diffusion.py
+__init__,,"def __init__(self, model, timestep_map, original_num_steps):
+    self.model = model
+    self.timestep_map = timestep_map
+    self.original_num_steps = original_num_steps
+",point_e\diffusion\gaussian_diffusion.py
+__call__,,"def __call__(self, x, ts, **kwargs):
+    map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
+    new_ts = map_tensor[ts]
+    return self.model(x, new_ts, **kwargs)
+",point_e\diffusion\gaussian_diffusion.py
+karras_sample,,"def karras_sample(*args, **kwargs):
+    last = None
+    for x in karras_sample_progressive(*args, **kwargs):
+        last = x['x']
+    return last
+",point_e\diffusion\k_diffusion.py
+karras_sample_progressive,,"def karras_sample_progressive(diffusion, model, shape, steps, clip_denoised
+    =True, progress=False, model_kwargs=None, device=None, sigma_min=0.002,
+    sigma_max=80, rho=7.0, sampler='heun', s_churn=0.0, s_tmin=0.0, s_tmax=
+    float('inf'), s_noise=1.0, guidance_scale=0.0):
+    sigmas = get_sigmas_karras(steps, sigma_min, sigma_max, rho, device=device)
+    x_T = th.randn(*shape, device=device) * sigma_max
+    sample_fn = {'heun': sample_heun, 'dpm': sample_dpm, 'ancestral':
+        sample_euler_ancestral}[sampler]
+    if sampler != 'ancestral':
+        sampler_args = dict(s_churn=s_churn, s_tmin=s_tmin, s_tmax=s_tmax,
+            s_noise=s_noise)
+    else:
+        sampler_args = {}
+    if isinstance(diffusion, KarrasDenoiser):
+
+        def denoiser(x_t, sigma):
+            _, denoised = diffusion.denoise(model, x_t, sigma, **model_kwargs)
+            if clip_denoised:
+                denoised = denoised.clamp(-1, 1)
+            return denoised
+    elif isinstance(diffusion, GaussianDiffusion):
+        model = GaussianToKarrasDenoiser(model, diffusion)
+
+        def denoiser(x_t, sigma):
+            _, denoised = model.denoise(x_t, sigma, clip_denoised=
+                clip_denoised, model_kwargs=model_kwargs)
+            return denoised
+    else:
+        raise NotImplementedError
+    if guidance_scale != 0 and guidance_scale != 1:
+
+        def guided_denoiser(x_t, sigma):
+            x_t = th.cat([x_t, x_t], dim=0)
+            sigma = th.cat([sigma, sigma], dim=0)
+            x_0 = denoiser(x_t, sigma)
+            cond_x_0, uncond_x_0 = th.split(x_0, len(x_0) // 2, dim=0)
+            x_0 = uncond_x_0 + guidance_scale * (cond_x_0 - uncond_x_0)
+            return x_0
+    else:
+        guided_denoiser = denoiser
+    for obj in sample_fn(guided_denoiser, x_T, sigmas, progress=progress,
+        **sampler_args):
+        if isinstance(diffusion, GaussianDiffusion):
+            yield diffusion.unscale_out_dict(obj)
+        else:
+            yield obj
+",point_e\diffusion\k_diffusion.py
+get_sigmas_karras,Constructs the noise schedule of Karras et al. (2022).,"def get_sigmas_karras(n, sigma_min, sigma_max, rho=7.0, device='cpu'):
+    """"""""""""
+    ramp = th.linspace(0, 1, n)
+    min_inv_rho = sigma_min ** (1 / rho)
+    max_inv_rho = sigma_max ** (1 / rho)
+    sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
+    return append_zero(sigmas).to(device)
+",point_e\diffusion\k_diffusion.py
+to_d,Converts a denoiser output to a Karras ODE derivative.,"def to_d(x, sigma, denoised):
+    """"""""""""
+    return (x - denoised) / append_dims(sigma, x.ndim)
+",point_e\diffusion\k_diffusion.py
+get_ancestral_step,"Calculates the noise level (sigma_down) to step down to and the amount
+of noise to add (sigma_up) when doing an ancestral sampling step.","def get_ancestral_step(sigma_from, sigma_to):
+    """"""""""""
+    sigma_up = (sigma_to ** 2 * (sigma_from ** 2 - sigma_to ** 2) / 
+        sigma_from ** 2) ** 0.5
+    sigma_down = (sigma_to ** 2 - sigma_up ** 2) ** 0.5
+    return sigma_down, sigma_up
+",point_e\diffusion\k_diffusion.py
+sample_euler_ancestral,Ancestral sampling with Euler method steps.,"@th.no_grad()
+def sample_euler_ancestral(model, x, sigmas, progress=False):
+    """"""""""""
+    s_in = x.new_ones([x.shape[0]])
+    indices = range(len(sigmas) - 1)
+    if progress:
+        from tqdm.auto import tqdm
+        indices = tqdm(indices)
+    for i in indices:
+        denoised = model(x, sigmas[i] * s_in)
+        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1])
+        yield {'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i],
+            'pred_xstart': denoised}
+        d = to_d(x, sigmas[i], denoised)
+        dt = sigma_down - sigmas[i]
+        x = x + d * dt
+        x = x + th.randn_like(x) * sigma_up
+    yield {'x': x, 'pred_xstart': x}
+",point_e\diffusion\k_diffusion.py
+sample_heun,Implements Algorithm 2 (Heun steps) from Karras et al. (2022).,"@th.no_grad()
+def sample_heun(denoiser, x, sigmas, progress=False, s_churn=0.0, s_tmin=
+    0.0, s_tmax=float('inf'), s_noise=1.0):
+    """"""""""""
+    s_in = x.new_ones([x.shape[0]])
+    indices = range(len(sigmas) - 1)
+    if progress:
+        from tqdm.auto import tqdm
+        indices = tqdm(indices)
+    for i in indices:
+        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1
+            ) if s_tmin <= sigmas[i] <= s_tmax else 0.0
+        eps = th.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+        denoised = denoiser(x, sigma_hat * s_in)
+        d = to_d(x, sigma_hat, denoised)
+        yield {'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat,
+            'pred_xstart': denoised}
+        dt = sigmas[i + 1] - sigma_hat
+        if sigmas[i + 1] == 0:
+            x = x + d * dt
+        else:
+            x_2 = x + d * dt
+            denoised_2 = denoiser(x_2, sigmas[i + 1] * s_in)
+            d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
+            d_prime = (d + d_2) / 2
+            x = x + d_prime * dt
+    yield {'x': x, 'pred_xstart': denoised}
+",point_e\diffusion\k_diffusion.py
+sample_dpm,A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022).,"@th.no_grad()
+def sample_dpm(denoiser, x, sigmas, progress=False, s_churn=0.0, s_tmin=0.0,
+    s_tmax=float('inf'), s_noise=1.0):
+    """"""""""""
+    s_in = x.new_ones([x.shape[0]])
+    indices = range(len(sigmas) - 1)
+    if progress:
+        from tqdm.auto import tqdm
+        indices = tqdm(indices)
+    for i in indices:
+        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1
+            ) if s_tmin <= sigmas[i] <= s_tmax else 0.0
+        eps = th.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+        denoised = denoiser(x, sigma_hat * s_in)
+        d = to_d(x, sigma_hat, denoised)
+        yield {'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat,
+            'denoised': denoised}
+        sigma_mid = ((sigma_hat ** (1 / 3) + sigmas[i + 1] ** (1 / 3)) / 2
+            ) ** 3
+        dt_1 = sigma_mid - sigma_hat
+        dt_2 = sigmas[i + 1] - sigma_hat
+        x_2 = x + d * dt_1
+        denoised_2 = denoiser(x_2, sigma_mid * s_in)
+        d_2 = to_d(x_2, sigma_mid, denoised_2)
+        x = x + d_2 * dt_2
+    yield {'x': x, 'pred_xstart': denoised}
+",point_e\diffusion\k_diffusion.py
+append_dims,Appends dimensions to the end of a tensor until it has target_dims dimensions.,"def append_dims(x, target_dims):
+    """"""""""""
+    dims_to_append = target_dims - x.ndim
+    if dims_to_append < 0:
+        raise ValueError(
+            f'input has {x.ndim} dims but target_dims is {target_dims}, which is less'
+            )
+    return x[(...,) + (None,) * dims_to_append]
+",point_e\diffusion\k_diffusion.py
+append_zero,,"def append_zero(x):
+    return th.cat([x, x.new_zeros([1])])
+",point_e\diffusion\k_diffusion.py
+__init__,,"def __init__(self, sigma_data: float=0.5):
+    self.sigma_data = sigma_data
+",point_e\diffusion\k_diffusion.py
+get_snr,,"def get_snr(self, sigmas):
+    return sigmas ** -2
+",point_e\diffusion\k_diffusion.py
+get_sigmas,,"def get_sigmas(self, sigmas):
+    return sigmas
+",point_e\diffusion\k_diffusion.py
+get_scalings,,"def get_scalings(self, sigma):
+    c_skip = self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2)
+    c_out = sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2
+        ) ** 0.5
+    c_in = 1 / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
+    return c_skip, c_out, c_in
+",point_e\diffusion\k_diffusion.py
+training_losses,,"def training_losses(self, model, x_start, sigmas, model_kwargs=None, noise=None
+    ):
+    if model_kwargs is None:
+        model_kwargs = {}
+    if noise is None:
+        noise = th.randn_like(x_start)
+    terms = {}
+    dims = x_start.ndim
+    x_t = x_start + noise * append_dims(sigmas, dims)
+    c_skip, c_out, _ = [append_dims(x, dims) for x in self.get_scalings(sigmas)
+        ]
+    model_output, denoised = self.denoise(model, x_t, sigmas, **model_kwargs)
+    target = (x_start - c_skip * x_t) / c_out
+    terms['mse'] = mean_flat((model_output - target) ** 2)
+    terms['xs_mse'] = mean_flat((denoised - x_start) ** 2)
+    if 'vb' in terms:
+        terms['loss'] = terms['mse'] + terms['vb']
+    else:
+        terms['loss'] = terms['mse']
+    return terms
+",point_e\diffusion\k_diffusion.py
+denoise,,"def denoise(self, model, x_t, sigmas, **model_kwargs):
+    c_skip, c_out, c_in = [append_dims(x, x_t.ndim) for x in self.
+        get_scalings(sigmas)]
+    rescaled_t = 1000 * 0.25 * th.log(sigmas + 1e-44)
+    model_output = model(c_in * x_t, rescaled_t, **model_kwargs)
+    denoised = c_out * model_output + c_skip * x_t
+    return model_output, denoised
+",point_e\diffusion\k_diffusion.py
+__init__,,"def __init__(self, model, diffusion):
+    from scipy import interpolate
+    self.model = model
+    self.diffusion = diffusion
+    self.alpha_cumprod_to_t = interpolate.interp1d(diffusion.alphas_cumprod,
+        np.arange(0, diffusion.num_timesteps))
+",point_e\diffusion\k_diffusion.py
+sigma_to_t,,"def sigma_to_t(self, sigma):
+    alpha_cumprod = 1.0 / (sigma ** 2 + 1)
+    if alpha_cumprod > self.diffusion.alphas_cumprod[0]:
+        return 0
+    elif alpha_cumprod <= self.diffusion.alphas_cumprod[-1]:
+        return self.diffusion.num_timesteps - 1
+    else:
+        return float(self.alpha_cumprod_to_t(alpha_cumprod))
+",point_e\diffusion\k_diffusion.py
+denoise,,"def denoise(self, x_t, sigmas, clip_denoised=True, model_kwargs=None):
+    t = th.tensor([self.sigma_to_t(sigma) for sigma in sigmas.cpu().numpy()
+        ], dtype=th.long, device=sigmas.device)
+    c_in = append_dims(1.0 / (sigmas ** 2 + 1) ** 0.5, x_t.ndim)
+    out = self.diffusion.p_mean_variance(self.model, x_t * c_in, t,
+        clip_denoised=clip_denoised, model_kwargs=model_kwargs)
+    return None, out['pred_xstart']
+",point_e\diffusion\k_diffusion.py
+__init__,,"def __init__(self, device: torch.device, models: Sequence[nn.Module],
+    diffusions: Sequence[GaussianDiffusion], num_points: Sequence[int],
+    aux_channels: Sequence[str], model_kwargs_key_filter: Sequence[str]=(
+    '*',), guidance_scale: Sequence[float]=(3.0, 3.0), clip_denoised: bool=
+    True, use_karras: Sequence[bool]=(True, True), karras_steps: Sequence[
+    int]=(64, 64), sigma_min: Sequence[float]=(0.001, 0.001), sigma_max:
+    Sequence[float]=(120, 160), s_churn: Sequence[float]=(3, 0)):
+    n = len(models)
+    assert n > 0
+    if n > 1:
+        if len(guidance_scale) == 1:
+            guidance_scale = list(guidance_scale) + [1.0] * (n - 1)
+        if len(use_karras) == 1:
+            use_karras = use_karras * n
+        if len(karras_steps) == 1:
+            karras_steps = karras_steps * n
+        if len(sigma_min) == 1:
+            sigma_min = sigma_min * n
+        if len(sigma_max) == 1:
+            sigma_max = sigma_max * n
+        if len(s_churn) == 1:
+            s_churn = s_churn * n
+        if len(model_kwargs_key_filter) == 1:
+            model_kwargs_key_filter = model_kwargs_key_filter * n
+    if len(model_kwargs_key_filter) == 0:
+        model_kwargs_key_filter = ['*'] * n
+    assert len(guidance_scale) == n
+    assert len(use_karras) == n
+    assert len(karras_steps) == n
+    assert len(sigma_min) == n
+    assert len(sigma_max) == n
+    assert len(s_churn) == n
+    assert len(model_kwargs_key_filter) == n
+    self.device = device
+    self.num_points = num_points
+    self.aux_channels = aux_channels
+    self.model_kwargs_key_filter = model_kwargs_key_filter
+    self.guidance_scale = guidance_scale
+    self.clip_denoised = clip_denoised
+    self.use_karras = use_karras
+    self.karras_steps = karras_steps
+    self.sigma_min = sigma_min
+    self.sigma_max = sigma_max
+    self.s_churn = s_churn
+    self.models = models
+    self.diffusions = diffusions
+",point_e\diffusion\sampler.py
+num_stages,,"@property
+def num_stages(self) ->int:
+    return len(self.models)
+",point_e\diffusion\sampler.py
+sample_batch,,"def sample_batch(self, batch_size: int, model_kwargs: Dict[str, Any]
+    ) ->torch.Tensor:
+    samples = None
+    for x in self.sample_batch_progressive(batch_size, model_kwargs):
+        samples = x
+    return samples
+",point_e\diffusion\sampler.py
+sample_batch_progressive,,"def sample_batch_progressive(self, batch_size: int, model_kwargs: Dict[str,
+    Any]) ->Iterator[torch.Tensor]:
+    samples = None
+    for model, diffusion, stage_num_points, stage_guidance_scale, stage_use_karras, stage_karras_steps, stage_sigma_min, stage_sigma_max, stage_s_churn, stage_key_filter in zip(
+        self.models, self.diffusions, self.num_points, self.guidance_scale,
+        self.use_karras, self.karras_steps, self.sigma_min, self.sigma_max,
+        self.s_churn, self.model_kwargs_key_filter):
+        stage_model_kwargs = model_kwargs.copy()
+        if stage_key_filter != '*':
+            use_keys = set(stage_key_filter.split(','))
+            stage_model_kwargs = {k: v for k, v in stage_model_kwargs.items
+                () if k in use_keys}
+        if samples is not None:
+            stage_model_kwargs['low_res'] = samples
+        if hasattr(model, 'cached_model_kwargs'):
+            stage_model_kwargs = model.cached_model_kwargs(batch_size,
+                stage_model_kwargs)
+        sample_shape = batch_size, 3 + len(self.aux_channels), stage_num_points
+        if stage_guidance_scale != 1 and stage_guidance_scale != 0:
+            for k, v in stage_model_kwargs.copy().items():
+                stage_model_kwargs[k] = torch.cat([v, torch.zeros_like(v)],
+                    dim=0)
+        if stage_use_karras:
+            samples_it = karras_sample_progressive(diffusion=diffusion,
+                model=model, shape=sample_shape, steps=stage_karras_steps,
+                clip_denoised=self.clip_denoised, model_kwargs=
+                stage_model_kwargs, device=self.device, sigma_min=
+                stage_sigma_min, sigma_max=stage_sigma_max, s_churn=
+                stage_s_churn, guidance_scale=stage_guidance_scale)
+        else:
+            internal_batch_size = batch_size
+            if stage_guidance_scale:
+                model = self._uncond_guide_model(model, stage_guidance_scale)
+                internal_batch_size *= 2
+            samples_it = diffusion.p_sample_loop_progressive(model, shape=(
+                internal_batch_size, *sample_shape[1:]), model_kwargs=
+                stage_model_kwargs, device=self.device, clip_denoised=self.
+                clip_denoised)
+        for x in samples_it:
+            samples = x['pred_xstart'][:batch_size]
+            if 'low_res' in stage_model_kwargs:
+                samples = torch.cat([stage_model_kwargs['low_res'][:len(
+                    samples)], samples], dim=-1)
+            yield samples
+",point_e\diffusion\sampler.py
+combine,,"@classmethod
+def combine(cls, *samplers: 'PointCloudSampler') ->'PointCloudSampler':
+    assert all(x.device == samplers[0].device for x in samplers[1:])
+    assert all(x.aux_channels == samplers[0].aux_channels for x in samplers[1:]
+        )
+    assert all(x.clip_denoised == samplers[0].clip_denoised for x in
+        samplers[1:])
+    return cls(device=samplers[0].device, models=[x for y in samplers for x in
+        y.models], diffusions=[x for y in samplers for x in y.diffusions],
+        num_points=[x for y in samplers for x in y.num_points],
+        aux_channels=samplers[0].aux_channels, model_kwargs_key_filter=[x for
+        y in samplers for x in y.model_kwargs_key_filter], guidance_scale=[
+        x for y in samplers for x in y.guidance_scale], clip_denoised=
+        samplers[0].clip_denoised, use_karras=[x for y in samplers for x in
+        y.use_karras], karras_steps=[x for y in samplers for x in y.
+        karras_steps], sigma_min=[x for y in samplers for x in y.sigma_min],
+        sigma_max=[x for y in samplers for x in y.sigma_max], s_churn=[x for
+        y in samplers for x in y.s_churn])
+",point_e\diffusion\sampler.py
+_uncond_guide_model,,"def _uncond_guide_model(self, model: Callable[..., torch.Tensor], scale: float
+    ) ->Callable[..., torch.Tensor]:
+
+    def model_fn(x_t, ts, **kwargs):
+        half = x_t[:len(x_t) // 2]
+        combined = torch.cat([half, half], dim=0)
+        model_out = model(combined, ts, **kwargs)
+        eps, rest = model_out[:, :3], model_out[:, 3:]
+        cond_eps, uncond_eps = torch.chunk(eps, 2, dim=0)
+        half_eps = uncond_eps + scale * (cond_eps - uncond_eps)
+        eps = torch.cat([half_eps, half_eps], dim=0)
+        return torch.cat([eps, rest], dim=1)
+    return model_fn
+",point_e\diffusion\sampler.py
+split_model_output,,"def split_model_output(self, output: torch.Tensor, rescale_colors: bool=False
+    ) ->Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+    assert len(self.aux_channels) + 3 == output.shape[1
+        ], 'there must be three spatial channels before aux'
+    pos, joined_aux = output[:, :3], output[:, 3:]
+    aux = {}
+    for i, name in enumerate(self.aux_channels):
+        v = joined_aux[:, i]
+        if name in {'R', 'G', 'B', 'A'}:
+            v = v.clamp(0, 255).round()
+            if rescale_colors:
+                v = v / 255.0
+        aux[name] = v
+    return pos, aux
+",point_e\diffusion\sampler.py
+output_to_point_clouds,,"def output_to_point_clouds(self, output: torch.Tensor) ->List[PointCloud]:
+    res = []
+    for sample in output:
+        xyz, aux = self.split_model_output(sample[None], rescale_colors=True)
+        res.append(PointCloud(coords=xyz[0].t().cpu().numpy(), channels={k:
+            v[0].cpu().numpy() for k, v in aux.items()}))
+    return res
+",point_e\diffusion\sampler.py
+with_options,,"def with_options(self, guidance_scale: float, clip_denoised: bool,
+    use_karras: Sequence[bool]=(True, True), karras_steps: Sequence[int]=(
+    64, 64), sigma_min: Sequence[float]=(0.001, 0.001), sigma_max: Sequence
+    [float]=(120, 160), s_churn: Sequence[float]=(3, 0)) ->'PointCloudSampler':
+    return PointCloudSampler(device=self.device, models=self.models,
+        diffusions=self.diffusions, num_points=self.num_points,
+        aux_channels=self.aux_channels, model_kwargs_key_filter=self.
+        model_kwargs_key_filter, guidance_scale=guidance_scale,
+        clip_denoised=clip_denoised, use_karras=use_karras, karras_steps=
+        karras_steps, sigma_min=sigma_min, sigma_max=sigma_max, s_churn=s_churn
+        )
+",point_e\diffusion\sampler.py
+get_torch_devices,,"def get_torch_devices() ->List[Union[str, torch.device]]:
+    if torch.cuda.is_available():
+        return [torch.device(f'cuda:{i}') for i in range(torch.cuda.
+            device_count())]
+    else:
+        return ['cpu']
+",point_e\evals\feature_extractor.py
+normalize_point_clouds,,"def normalize_point_clouds(pc: np.ndarray) ->np.ndarray:
+    centroids = np.mean(pc, axis=1, keepdims=True)
+    pc = pc - centroids
+    m = np.max(np.sqrt(np.sum(pc ** 2, axis=-1, keepdims=True)), axis=1,
+        keepdims=True)
+    pc = pc / m
+    return pc
+",point_e\evals\feature_extractor.py
+supports_predictions,,"@property
+@abstractmethod
+def supports_predictions(self) ->bool:
+    pass
+",point_e\evals\feature_extractor.py
+feature_dim,,"@property
+@abstractmethod
+def feature_dim(self) ->int:
+    pass
+",point_e\evals\feature_extractor.py
+num_classes,,"@property
+@abstractmethod
+def num_classes(self) ->int:
+    pass
+",point_e\evals\feature_extractor.py
+features_and_preds,"For a stream of point cloud batches, compute feature vectors and class
+predictions.
+
+:param point_clouds: a streamer for a sample batch. Typically, arr_0
+                     will contain the XYZ coordinates.
+:return: a tuple (features, predictions)
+         - features: a [B x feature_dim] array of feature vectors.
+         - predictions: a [B x num_classes] array of probabilities.","@abstractmethod
+def features_and_preds(self, streamer: NpzStreamer) ->Tuple[np.ndarray, np.
+    ndarray]:
+    """"""""""""
+",point_e\evals\feature_extractor.py
+__init__,,"def __init__(self, devices: List[Union[str, torch.device]],
+    device_batch_size: int=64, cache_dir: Optional[str]=None):
+    state_dict = load_checkpoint('pointnet', device=torch.device('cpu'),
+        cache_dir=cache_dir)['model_state_dict']
+    self.device_batch_size = device_batch_size
+    self.devices = devices
+    self.models = []
+    for device in devices:
+        model = get_model(num_class=40, normal_channel=False, width_mult=2)
+        model.load_state_dict(state_dict)
+        model.to(device)
+        model.eval()
+        self.models.append(model)
+",point_e\evals\feature_extractor.py
+supports_predictions,,"@property
+def supports_predictions(self) ->bool:
+    return True
+",point_e\evals\feature_extractor.py
+feature_dim,,"@property
+def feature_dim(self) ->int:
+    return 256
+",point_e\evals\feature_extractor.py
+num_classes,,"@property
+def num_classes(self) ->int:
+    return 40
+",point_e\evals\feature_extractor.py
+features_and_preds,,"def features_and_preds(self, streamer: NpzStreamer) ->Tuple[np.ndarray, np.
+    ndarray]:
+    batch_size = self.device_batch_size * len(self.devices)
+    point_clouds = (x['arr_0'] for x in streamer.stream(batch_size, ['arr_0']))
+    output_features = []
+    output_predictions = []
+    with ThreadPool(len(self.devices)) as pool:
+        for batch in point_clouds:
+            batch = normalize_point_clouds(batch)
+            batches = []
+            for i, device in zip(range(0, len(batch), self.
+                device_batch_size), self.devices):
+                batches.append(torch.from_numpy(batch[i:i + self.
+                    device_batch_size]).permute(0, 2, 1).to(dtype=torch.
+                    float32, device=device))
+
+            def compute_features(i_batch):
+                i, batch = i_batch
+                with torch.no_grad():
+                    return self.models[i](batch, features=True)
+            for logits, _, features in pool.imap(compute_features,
+                enumerate(batches)):
+                output_features.append(features.cpu().numpy())
+                output_predictions.append(logits.exp().cpu().numpy())
+    return np.concatenate(output_features, axis=0), np.concatenate(
+        output_predictions, axis=0)
+",point_e\evals\feature_extractor.py
+compute_statistics,,"def compute_statistics(feats: np.ndarray) ->FIDStatistics:
+    mu = np.mean(feats, axis=0)
+    sigma = np.cov(feats, rowvar=False)
+    return FIDStatistics(mu, sigma)
+",point_e\evals\fid_is.py
+compute_inception_score,,"def compute_inception_score(preds: np.ndarray, split_size: int=5000) ->float:
+    scores = []
+    for i in range(0, len(preds), split_size):
+        part = preds[i:i + split_size]
+        kl = part * (np.log(part) - np.log(np.expand_dims(np.mean(part, 0), 0))
+            )
+        kl = np.mean(np.sum(kl, 1))
+        scores.append(np.exp(kl))
+    return float(np.mean(scores))
+",point_e\evals\fid_is.py
+__init__,,"def __init__(self, mu: np.ndarray, sigma: np.ndarray):
+    self.mu = mu
+    self.sigma = sigma
+",point_e\evals\fid_is.py
+frechet_distance,Compute the Frechet distance between two sets of statistics.,"def frechet_distance(self, other, eps=1e-06):
+    """"""""""""
+    mu1, sigma1 = self.mu, self.sigma
+    mu2, sigma2 = other.mu, other.sigma
+    mu1 = np.atleast_1d(mu1)
+    mu2 = np.atleast_1d(mu2)
+    sigma1 = np.atleast_2d(sigma1)
+    sigma2 = np.atleast_2d(sigma2)
+    assert mu1.shape == mu2.shape, f'Training and test mean vectors have different lengths: {mu1.shape}, {mu2.shape}'
+    assert sigma1.shape == sigma2.shape, f'Training and test covariances have different dimensions: {sigma1.shape}, {sigma2.shape}'
+    diff = mu1 - mu2
+    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+    if not np.isfinite(covmean).all():
+        msg = (
+            'fid calculation produces singular product; adding %s to diagonal of cov estimates'
+             % eps)
+        warnings.warn(msg)
+        offset = np.eye(sigma1.shape[0]) * eps
+        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+    if np.iscomplexobj(covmean):
+        if not np.allclose(np.diagonal(covmean).imag, 0, atol=0.001):
+            m = np.max(np.abs(covmean.imag))
+            raise ValueError('Imaginary component {}'.format(m))
+        covmean = covmean.real
+    tr_covmean = np.trace(covmean)
+    return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2
+        ) - 2 * tr_covmean
+",point_e\evals\fid_is.py
+_npz_paths_and_length,,"def _npz_paths_and_length(glob_path: str) ->Tuple[List[str], Optional[int]]:
+    count_match = re.match('^(.*)\\[:([0-9]*)\\]$', glob_path)
+    if count_match:
+        raw_path = count_match[1]
+        max_count = int(count_match[2])
+    else:
+        raw_path = glob_path
+        max_count = None
+    paths = sorted(glob.glob(raw_path))
+    if not len(paths):
+        raise ValueError(f'no paths found matching: {glob_path}')
+    return paths, max_count
+",point_e\evals\npz_stream.py
+open_npz_arrays,,"@contextmanager
+def open_npz_arrays(path: str, arr_names: Sequence[str]) ->List[NpzArrayReader
+    ]:
+    if not len(arr_names):
+        yield []
+        return
+    arr_name = arr_names[0]
+    with open_array(path, arr_name) as arr_f:
+        version = np.lib.format.read_magic(arr_f)
+        header = None
+        if version == (1, 0):
+            header = np.lib.format.read_array_header_1_0(arr_f)
+        elif version == (2, 0):
+            header = np.lib.format.read_array_header_2_0(arr_f)
+        if header is None:
+            reader = MemoryNpzArrayReader.load(path, arr_name)
+        else:
+            shape, fortran, dtype = header
+            if fortran or dtype.hasobject:
+                reader = MemoryNpzArrayReader.load(path, arr_name)
+            else:
+                reader = StreamingNpzArrayReader(arr_f, shape, dtype)
+        with open_npz_arrays(path, arr_names[1:]) as next_readers:
+            yield [reader] + next_readers
+",point_e\evals\npz_stream.py
+_read_bytes,"Copied from: https://github.com/numpy/numpy/blob/fb215c76967739268de71aa4bda55dd1b062bc2e/numpy/lib/format.py#L788-L886
+
+Read from file-like object until size bytes are read.
+Raises ValueError if not EOF is encountered before size bytes are read.
+Non-blocking objects only supported if they derive from io objects.
+Required as e.g. ZipExtFile in python 2.6 can return less data than
+requested.","def _read_bytes(fp, size, error_template='ran out of data'):
+    """"""""""""
+    data = bytes()
+    while True:
+        try:
+            r = fp.read(size - len(data))
+            data += r
+            if len(r) == 0 or len(data) == size:
+                break
+        except io.BlockingIOError:
+            pass
+    if len(data) != size:
+        msg = 'EOF: reading %s, expected %d bytes got %d'
+        raise ValueError(msg % (error_template, size, len(data)))
+    else:
+        return data
+",point_e\evals\npz_stream.py
+open_array,,"@contextmanager
+def open_array(path: str, arr_name: str):
+    with open(path, 'rb') as f:
+        with zipfile.ZipFile(f, 'r') as zip_f:
+            if f'{arr_name}.npy' not in zip_f.namelist():
+                raise ValueError(f'missing {arr_name} in npz file')
+            with zip_f.open(f'{arr_name}.npy', 'r') as arr_f:
+                yield arr_f
+",point_e\evals\npz_stream.py
+_dict_batch_size,,"def _dict_batch_size(objs: Dict[str, np.ndarray]) ->int:
+    return len(next(iter(objs.values())))
+",point_e\evals\npz_stream.py
+infos_from_first_file,,"@classmethod
+def infos_from_first_file(cls, glob_path: str) ->Dict[str, 'NumpyArrayInfo']:
+    paths, _ = _npz_paths_and_length(glob_path)
+    return cls.infos_from_file(paths[0])
+",point_e\evals\npz_stream.py
+infos_from_file,Extract the info of every array in an npz file.,"@classmethod
+def infos_from_file(cls, npz_path: str) ->Dict[str, 'NumpyArrayInfo']:
+    """"""""""""
+    if not os.path.exists(npz_path):
+        raise FileNotFoundError(f'batch of samples was not found: {npz_path}')
+    results = {}
+    with open(npz_path, 'rb') as f:
+        with zipfile.ZipFile(f, 'r') as zip_f:
+            for name in zip_f.namelist():
+                if not name.endswith('.npy'):
+                    continue
+                key_name = name[:-len('.npy')]
+                with zip_f.open(name, 'r') as arr_f:
+                    version = np.lib.format.read_magic(arr_f)
+                    if version == (1, 0):
+                        header = np.lib.format.read_array_header_1_0(arr_f)
+                    elif version == (2, 0):
+                        header = np.lib.format.read_array_header_2_0(arr_f)
+                    else:
+                        raise ValueError(
+                            f'unknown numpy array version: {version}')
+                    shape, _, dtype = header
+                    results[key_name] = cls(name=key_name, dtype=dtype,
+                        shape=shape)
+    return results
+",point_e\evals\npz_stream.py
+elem_shape,,"@property
+def elem_shape(self) ->Tuple[int]:
+    return self.shape[1:]
+",point_e\evals\npz_stream.py
+validate,,"def validate(self):
+    if self.name in {'R', 'G', 'B'}:
+        if len(self.shape) != 2:
+            raise ValueError(
+                f""expecting exactly 2-D shape for '{self.name}' but got: {self.shape}""
+                )
+    elif self.name == 'arr_0':
+        if len(self.shape) < 2:
+            raise ValueError(
+                f'expecting at least 2-D shape but got: {self.shape}')
+        elif len(self.shape) == 3:
+            if not np.issubdtype(self.dtype, np.floating):
+                raise ValueError(
+                    f'invalid dtype for audio batch: {self.dtype} (expected float)'
+                    )
+        elif self.dtype != np.uint8:
+            raise ValueError(
+                f'invalid dtype for image batch: {self.dtype} (expected uint8)'
+                )
+",point_e\evals\npz_stream.py
+__init__,,"def __init__(self, glob_path: str):
+    self.paths, self.trunc_length = _npz_paths_and_length(glob_path)
+    self.infos = NumpyArrayInfo.infos_from_file(self.paths[0])
+",point_e\evals\npz_stream.py
+keys,,"def keys(self) ->List[str]:
+    return list(self.infos.keys())
+",point_e\evals\npz_stream.py
+stream,,"def stream(self, batch_size: int, keys: Sequence[str]) ->Iterator[Dict[str,
+    np.ndarray]]:
+    cur_batch = None
+    num_remaining = self.trunc_length
+    for path in self.paths:
+        if num_remaining is not None and num_remaining <= 0:
+            break
+        with open_npz_arrays(path, keys) as readers:
+            combined_reader = CombinedReader(keys, readers)
+            while num_remaining is None or num_remaining > 0:
+                read_bs = batch_size
+                if cur_batch is not None:
+                    read_bs -= _dict_batch_size(cur_batch)
+                if num_remaining is not None:
+                    read_bs = min(read_bs, num_remaining)
+                batch = combined_reader.read_batch(read_bs)
+                if batch is None:
+                    break
+                if num_remaining is not None:
+                    num_remaining -= _dict_batch_size(batch)
+                if cur_batch is None:
+                    cur_batch = batch
+                else:
+                    cur_batch = {k: np.concatenate([cur_batch[k], v], axis=
+                        0) for k, v in batch.items()}
+                if _dict_batch_size(cur_batch) == batch_size:
+                    yield cur_batch
+                    cur_batch = None
+    if cur_batch is not None:
+        yield cur_batch
+",point_e\evals\npz_stream.py
+read_batch,,"@abstractmethod
+def read_batch(self, batch_size: int) ->Optional[np.ndarray]:
+    pass
+",point_e\evals\npz_stream.py
+__init__,,"def __init__(self, arr_f, shape, dtype):
+    self.arr_f = arr_f
+    self.shape = shape
+    self.dtype = dtype
+    self.idx = 0
+",point_e\evals\npz_stream.py
+read_batch,,"def read_batch(self, batch_size: int) ->Optional[np.ndarray]:
+    if self.idx >= self.shape[0]:
+        return None
+    bs = min(batch_size, self.shape[0] - self.idx)
+    self.idx += bs
+    if self.dtype.itemsize == 0:
+        return np.ndarray([bs, *self.shape[1:]], dtype=self.dtype)
+    read_count = bs * np.prod(self.shape[1:])
+    read_size = int(read_count * self.dtype.itemsize)
+    data = _read_bytes(self.arr_f, read_size, 'array data')
+    return np.frombuffer(data, dtype=self.dtype).reshape([bs, *self.shape[1:]])
+",point_e\evals\npz_stream.py
+__init__,,"def __init__(self, arr):
+    self.arr = arr
+    self.idx = 0
+",point_e\evals\npz_stream.py
+load,,"@classmethod
+def load(cls, path: str, arr_name: str):
+    with open(path, 'rb') as f:
+        arr = np.load(f)[arr_name]
+    return cls(arr)
+",point_e\evals\npz_stream.py
+read_batch,,"def read_batch(self, batch_size: int) ->Optional[np.ndarray]:
+    if self.idx >= self.arr.shape[0]:
+        return None
+    res = self.arr[self.idx:self.idx + batch_size]
+    self.idx += batch_size
+    return res
+",point_e\evals\npz_stream.py
+__init__,,"def __init__(self, keys: List[str], readers: List[NpzArrayReader]):
+    self.keys = keys
+    self.readers = readers
+",point_e\evals\npz_stream.py
+read_batch,,"def read_batch(self, batch_size: int) ->Optional[Dict[str, np.ndarray]]:
+    batches = [r.read_batch(batch_size) for r in self.readers]
+    any_none = any(x is None for x in batches)
+    all_none = all(x is None for x in batches)
+    if any_none != all_none:
+        raise RuntimeError('different keys had different numbers of elements')
+    if any_none:
+        return None
+    if any(len(x) != len(batches[0]) for x in batches):
+        raise RuntimeError('different keys had different numbers of elements')
+    return dict(zip(self.keys, batches))
+",point_e\evals\npz_stream.py
+__init__,,"def __init__(self, num_class, normal_channel=True, width_mult=1):
+    super(get_model, self).__init__()
+    self.width_mult = width_mult
+    in_channel = 6 if normal_channel else 3
+    self.normal_channel = normal_channel
+    self.sa1 = PointNetSetAbstraction(npoint=512, radius=0.2, nsample=32,
+        in_channel=in_channel, mlp=[64 * width_mult, 64 * width_mult, 128 *
+        width_mult], group_all=False)
+    self.sa2 = PointNetSetAbstraction(npoint=128, radius=0.4, nsample=64,
+        in_channel=128 * width_mult + 3, mlp=[128 * width_mult, 128 *
+        width_mult, 256 * width_mult], group_all=False)
+    self.sa3 = PointNetSetAbstraction(npoint=None, radius=None, nsample=
+        None, in_channel=256 * width_mult + 3, mlp=[256 * width_mult, 512 *
+        width_mult, 1024 * width_mult], group_all=True)
+    self.fc1 = nn.Linear(1024 * width_mult, 512 * width_mult)
+    self.bn1 = nn.BatchNorm1d(512 * width_mult)
+    self.drop1 = nn.Dropout(0.4)
+    self.fc2 = nn.Linear(512 * width_mult, 256 * width_mult)
+    self.bn2 = nn.BatchNorm1d(256 * width_mult)
+    self.drop2 = nn.Dropout(0.4)
+    self.fc3 = nn.Linear(256 * width_mult, num_class)
+",point_e\evals\pointnet2_cls_ssg.py
+forward,,"def forward(self, xyz, features=False):
+    B, _, _ = xyz.shape
+    if self.normal_channel:
+        norm = xyz[:, 3:, :]
+        xyz = xyz[:, :3, :]
+    else:
+        norm = None
+    l1_xyz, l1_points = self.sa1(xyz, norm)
+    l2_xyz, l2_points = self.sa2(l1_xyz, l1_points)
+    l3_xyz, l3_points = self.sa3(l2_xyz, l2_points)
+    x = l3_points.view(B, 1024 * self.width_mult)
+    x = self.drop1(F.relu(self.bn1(self.fc1(x))))
+    result_features = self.bn2(self.fc2(x))
+    x = self.drop2(F.relu(result_features))
+    x = self.fc3(x)
+    x = F.log_softmax(x, -1)
+    if features:
+        return x, l3_points, result_features
+    else:
+        return x, l3_points
+",point_e\evals\pointnet2_cls_ssg.py
+__init__,,"def __init__(self):
+    super(get_loss, self).__init__()
+",point_e\evals\pointnet2_cls_ssg.py
+forward,,"def forward(self, pred, target, trans_feat):
+    total_loss = F.nll_loss(pred, target)
+    return total_loss
+",point_e\evals\pointnet2_cls_ssg.py
+timeit,,"def timeit(tag, t):
+    print('{}: {}s'.format(tag, time() - t))
+    return time()
+",point_e\evals\pointnet2_utils.py
+pc_normalize,,"def pc_normalize(pc):
+    l = pc.shape[0]
+    centroid = np.mean(pc, axis=0)
+    pc = pc - centroid
+    m = np.max(np.sqrt(np.sum(pc ** 2, axis=1)))
+    pc = pc / m
+    return pc
+",point_e\evals\pointnet2_utils.py
+square_distance,"Calculate Euclid distance between each two points.
+
+src^T * dst = xn * xm + yn * ym + zn * zm;
+sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
+sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
+dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
+     = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst
+
+Input:
+    src: source points, [B, N, C]
+    dst: target points, [B, M, C]
+Output:
+    dist: per-point square distance, [B, N, M]","def square_distance(src, dst):
+    """"""""""""
+    B, N, _ = src.shape
+    _, M, _ = dst.shape
+    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
+    dist += torch.sum(src ** 2, -1).view(B, N, 1)
+    dist += torch.sum(dst ** 2, -1).view(B, 1, M)
+    return dist
+",point_e\evals\pointnet2_utils.py
+index_points,"Input:
+    points: input points data, [B, N, C]
+    idx: sample index data, [B, S]
+Return:
+    new_points:, indexed points data, [B, S, C]","def index_points(points, idx):
+    """"""""""""
+    device = points.device
+    B = points.shape[0]
+    view_shape = list(idx.shape)
+    view_shape[1:] = [1] * (len(view_shape) - 1)
+    repeat_shape = list(idx.shape)
+    repeat_shape[0] = 1
+    batch_indices = torch.arange(B, dtype=torch.long).to(device).view(
+        view_shape).repeat(repeat_shape)
+    new_points = points[batch_indices, idx, :]
+    return new_points
+",point_e\evals\pointnet2_utils.py
+farthest_point_sample,"Input:
+    xyz: pointcloud data, [B, N, 3]
+    npoint: number of samples
+Return:
+    centroids: sampled pointcloud index, [B, npoint]","def farthest_point_sample(xyz, npoint, deterministic=False):
+    """"""""""""
+    device = xyz.device
+    B, N, C = xyz.shape
+    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
+    distance = torch.ones(B, N).to(device) * 10000000000.0
+    if deterministic:
+        farthest = torch.arange(0, B, dtype=torch.long).to(device)
+    else:
+        farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device)
+    batch_indices = torch.arange(B, dtype=torch.long).to(device)
+    for i in range(npoint):
+        centroids[:, i] = farthest
+        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
+        dist = torch.sum((xyz - centroid) ** 2, -1)
+        mask = dist < distance
+        distance[mask] = dist[mask]
+        farthest = torch.max(distance, -1)[1]
+    return centroids
+",point_e\evals\pointnet2_utils.py
+query_ball_point,"Input:
+    radius: local region radius
+    nsample: max sample number in local region
+    xyz: all points, [B, N, 3]
+    new_xyz: query points, [B, S, 3]
+Return:
+    group_idx: grouped points index, [B, S, nsample]","def query_ball_point(radius, nsample, xyz, new_xyz):
+    """"""""""""
+    device = xyz.device
+    B, N, C = xyz.shape
+    _, S, _ = new_xyz.shape
+    group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N
+        ).repeat([B, S, 1])
+    sqrdists = square_distance(new_xyz, xyz)
+    group_idx[sqrdists > radius ** 2] = N
+    group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]
+    group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample])
+    mask = group_idx == N
+    group_idx[mask] = group_first[mask]
+    return group_idx
+",point_e\evals\pointnet2_utils.py
+sample_and_group,"Input:
+    npoint:
+    radius:
+    nsample:
+    xyz: input points position data, [B, N, 3]
+    points: input points data, [B, N, D]
+Return:
+    new_xyz: sampled points position data, [B, npoint, nsample, 3]
+    new_points: sampled points data, [B, npoint, nsample, 3+D]","def sample_and_group(npoint, radius, nsample, xyz, points, returnfps=False,
+    deterministic=False):
+    """"""""""""
+    B, N, C = xyz.shape
+    S = npoint
+    fps_idx = farthest_point_sample(xyz, npoint, deterministic=deterministic)
+    new_xyz = index_points(xyz, fps_idx)
+    idx = query_ball_point(radius, nsample, xyz, new_xyz)
+    grouped_xyz = index_points(xyz, idx)
+    grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C)
+    if points is not None:
+        grouped_points = index_points(points, idx)
+        new_points = torch.cat([grouped_xyz_norm, grouped_points], dim=-1)
+    else:
+        new_points = grouped_xyz_norm
+    if returnfps:
+        return new_xyz, new_points, grouped_xyz, fps_idx
+    else:
+        return new_xyz, new_points
+",point_e\evals\pointnet2_utils.py
+sample_and_group_all,"Input:
+    xyz: input points position data, [B, N, 3]
+    points: input points data, [B, N, D]
+Return:
+    new_xyz: sampled points position data, [B, 1, 3]
+    new_points: sampled points data, [B, 1, N, 3+D]","def sample_and_group_all(xyz, points):
+    """"""""""""
+    device = xyz.device
+    B, N, C = xyz.shape
+    new_xyz = torch.zeros(B, 1, C).to(device)
+    grouped_xyz = xyz.view(B, 1, N, C)
+    if points is not None:
+        new_points = torch.cat([grouped_xyz, points.view(B, 1, N, -1)], dim=-1)
+    else:
+        new_points = grouped_xyz
+    return new_xyz, new_points
+",point_e\evals\pointnet2_utils.py
+__init__,,"def __init__(self, npoint, radius, nsample, in_channel, mlp, group_all):
+    super(PointNetSetAbstraction, self).__init__()
+    self.npoint = npoint
+    self.radius = radius
+    self.nsample = nsample
+    self.mlp_convs = nn.ModuleList()
+    self.mlp_bns = nn.ModuleList()
+    last_channel = in_channel
+    for out_channel in mlp:
+        self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
+        self.mlp_bns.append(nn.BatchNorm2d(out_channel))
+        last_channel = out_channel
+    self.group_all = group_all
+",point_e\evals\pointnet2_utils.py
+forward,"Input:
+    xyz: input points position data, [B, C, N]
+    points: input points data, [B, D, N]
+Return:
+    new_xyz: sampled points position data, [B, C, S]
+    new_points_concat: sample points feature data, [B, D', S]","def forward(self, xyz, points):
+    """"""""""""
+    xyz = xyz.permute(0, 2, 1)
+    if points is not None:
+        points = points.permute(0, 2, 1)
+    if self.group_all:
+        new_xyz, new_points = sample_and_group_all(xyz, points)
+    else:
+        new_xyz, new_points = sample_and_group(self.npoint, self.radius,
+            self.nsample, xyz, points, deterministic=not self.training)
+    new_points = new_points.permute(0, 3, 2, 1)
+    for i, conv in enumerate(self.mlp_convs):
+        bn = self.mlp_bns[i]
+        new_points = F.relu(bn(conv(new_points)))
+    new_points = torch.max(new_points, 2)[0]
+    new_xyz = new_xyz.permute(0, 2, 1)
+    return new_xyz, new_points
+",point_e\evals\pointnet2_utils.py
+__init__,,"def __init__(self, npoint, radius_list, nsample_list, in_channel, mlp_list):
+    super(PointNetSetAbstractionMsg, self).__init__()
+    self.npoint = npoint
+    self.radius_list = radius_list
+    self.nsample_list = nsample_list
+    self.conv_blocks = nn.ModuleList()
+    self.bn_blocks = nn.ModuleList()
+    for i in range(len(mlp_list)):
+        convs = nn.ModuleList()
+        bns = nn.ModuleList()
+        last_channel = in_channel + 3
+        for out_channel in mlp_list[i]:
+            convs.append(nn.Conv2d(last_channel, out_channel, 1))
+            bns.append(nn.BatchNorm2d(out_channel))
+            last_channel = out_channel
+        self.conv_blocks.append(convs)
+        self.bn_blocks.append(bns)
+",point_e\evals\pointnet2_utils.py
+forward,"Input:
+    xyz: input points position data, [B, C, N]
+    points: input points data, [B, D, N]
+Return:
+    new_xyz: sampled points position data, [B, C, S]
+    new_points_concat: sample points feature data, [B, D', S]","def forward(self, xyz, points):
+    """"""""""""
+    xyz = xyz.permute(0, 2, 1)
+    if points is not None:
+        points = points.permute(0, 2, 1)
+    B, N, C = xyz.shape
+    S = self.npoint
+    new_xyz = index_points(xyz, farthest_point_sample(xyz, S, deterministic
+        =not self.training))
+    new_points_list = []
+    for i, radius in enumerate(self.radius_list):
+        K = self.nsample_list[i]
+        group_idx = query_ball_point(radius, K, xyz, new_xyz)
+        grouped_xyz = index_points(xyz, group_idx)
+        grouped_xyz -= new_xyz.view(B, S, 1, C)
+        if points is not None:
+            grouped_points = index_points(points, group_idx)
+            grouped_points = torch.cat([grouped_points, grouped_xyz], dim=-1)
+        else:
+            grouped_points = grouped_xyz
+        grouped_points = grouped_points.permute(0, 3, 2, 1)
+        for j in range(len(self.conv_blocks[i])):
+            conv = self.conv_blocks[i][j]
+            bn = self.bn_blocks[i][j]
+            grouped_points = F.relu(bn(conv(grouped_points)))
+        new_points = torch.max(grouped_points, 2)[0]
+        new_points_list.append(new_points)
+    new_xyz = new_xyz.permute(0, 2, 1)
+    new_points_concat = torch.cat(new_points_list, dim=1)
+    return new_xyz, new_points_concat
+",point_e\evals\pointnet2_utils.py
+__init__,,"def __init__(self, in_channel, mlp):
+    super(PointNetFeaturePropagation, self).__init__()
+    self.mlp_convs = nn.ModuleList()
+    self.mlp_bns = nn.ModuleList()
+    last_channel = in_channel
+    for out_channel in mlp:
+        self.mlp_convs.append(nn.Conv1d(last_channel, out_channel, 1))
+        self.mlp_bns.append(nn.BatchNorm1d(out_channel))
+        last_channel = out_channel
+",point_e\evals\pointnet2_utils.py
+forward,"Input:
+    xyz1: input points position data, [B, C, N]
+    xyz2: sampled input points position data, [B, C, S]
+    points1: input points data, [B, D, N]
+    points2: input points data, [B, D, S]
+Return:
+    new_points: upsampled points data, [B, D', N]","def forward(self, xyz1, xyz2, points1, points2):
+    """"""""""""
+    xyz1 = xyz1.permute(0, 2, 1)
+    xyz2 = xyz2.permute(0, 2, 1)
+    points2 = points2.permute(0, 2, 1)
+    B, N, C = xyz1.shape
+    _, S, _ = xyz2.shape
+    if S == 1:
+        interpolated_points = points2.repeat(1, N, 1)
+    else:
+        dists = square_distance(xyz1, xyz2)
+        dists, idx = dists.sort(dim=-1)
+        dists, idx = dists[:, :, :3], idx[:, :, :3]
+        dist_recip = 1.0 / (dists + 1e-08)
+        norm = torch.sum(dist_recip, dim=2, keepdim=True)
+        weight = dist_recip / norm
+        interpolated_points = torch.sum(index_points(points2, idx) * weight
+            .view(B, N, 3, 1), dim=2)
+    if points1 is not None:
+        points1 = points1.permute(0, 2, 1)
+        new_points = torch.cat([points1, interpolated_points], dim=-1)
+    else:
+        new_points = interpolated_points
+    new_points = new_points.permute(0, 2, 1)
+    for i, conv in enumerate(self.mlp_convs):
+        bn = self.mlp_bns[i]
+        new_points = F.relu(bn(conv(new_points)))
+    return new_points
+",point_e\evals\pointnet2_utils.py
+clear_scene,,"def clear_scene():
+    bpy.ops.object.select_all(action='SELECT')
+    bpy.ops.object.delete()
+",point_e\evals\scripts\blender_script.py
+clear_lights,,"def clear_lights():
+    bpy.ops.object.select_all(action='DESELECT')
+    for obj in bpy.context.scene.objects.values():
+        if isinstance(obj.data, bpy.types.Light):
+            obj.select_set(True)
+    bpy.ops.object.delete()
+",point_e\evals\scripts\blender_script.py
+import_model,,"def import_model(path):
+    clear_scene()
+    _, ext = os.path.splitext(path)
+    ext = ext.lower()
+    if ext == '.obj':
+        bpy.ops.import_scene.obj(filepath=path)
+    elif ext in ['.glb', '.gltf']:
+        bpy.ops.import_scene.gltf(filepath=path)
+    elif ext == '.stl':
+        bpy.ops.import_mesh.stl(filepath=path)
+    elif ext == '.fbx':
+        bpy.ops.import_scene.fbx(filepath=path)
+    elif ext == '.dae':
+        bpy.ops.wm.collada_import(filepath=path)
+    elif ext == '.ply':
+        bpy.ops.import_mesh.ply(filepath=path)
+    else:
+        raise RuntimeError(f'unexpected extension: {ext}')
+",point_e\evals\scripts\blender_script.py
+scene_root_objects,,"def scene_root_objects():
+    for obj in bpy.context.scene.objects.values():
+        if not obj.parent:
+            yield obj
+",point_e\evals\scripts\blender_script.py
+scene_bbox,,"def scene_bbox(single_obj=None, ignore_matrix=False):
+    bbox_min = (math.inf,) * 3
+    bbox_max = (-math.inf,) * 3
+    found = False
+    for obj in (scene_meshes() if single_obj is None else [single_obj]):
+        found = True
+        for coord in obj.bound_box:
+            coord = Vector(coord)
+            if not ignore_matrix:
+                coord = obj.matrix_world @ coord
+            bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
+            bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
+    if not found:
+        raise RuntimeError('no objects in scene to compute bounding box for')
+    return Vector(bbox_min), Vector(bbox_max)
+",point_e\evals\scripts\blender_script.py
+scene_meshes,,"def scene_meshes():
+    for obj in bpy.context.scene.objects.values():
+        if isinstance(obj.data, bpy.types.Mesh):
+            yield obj
+",point_e\evals\scripts\blender_script.py
+normalize_scene,,"def normalize_scene():
+    bbox_min, bbox_max = scene_bbox()
+    scale = 1 / max(bbox_max - bbox_min)
+    for obj in scene_root_objects():
+        obj.scale = obj.scale * scale
+    bpy.context.view_layer.update()
+    bbox_min, bbox_max = scene_bbox()
+    offset = -(bbox_min + bbox_max) / 2
+    for obj in scene_root_objects():
+        obj.matrix_world.translation += offset
+    bpy.ops.object.select_all(action='DESELECT')
+",point_e\evals\scripts\blender_script.py
+create_camera,,"def create_camera():
+    camera_data = bpy.data.cameras.new(name='Camera')
+    camera_object = bpy.data.objects.new('Camera', camera_data)
+    bpy.context.scene.collection.objects.link(camera_object)
+    bpy.context.scene.camera = camera_object
+",point_e\evals\scripts\blender_script.py
+set_camera,,"def set_camera(direction, camera_dist=2.0):
+    camera_pos = -camera_dist * direction
+    bpy.context.scene.camera.location = camera_pos
+    rot_quat = direction.to_track_quat('-Z', 'Y')
+    bpy.context.scene.camera.rotation_euler = rot_quat.to_euler()
+    bpy.context.view_layer.update()
+",point_e\evals\scripts\blender_script.py
+randomize_camera,,"def randomize_camera(camera_dist=2.0):
+    direction = random_unit_vector()
+    set_camera(direction, camera_dist=camera_dist)
+",point_e\evals\scripts\blender_script.py
+pan_camera,,"def pan_camera(time, axis='Z', camera_dist=2.0, elevation=-0.1):
+    angle = time * math.pi * 2
+    direction = [-math.cos(angle), -math.sin(angle), -elevation]
+    assert axis in ['X', 'Y', 'Z']
+    if axis == 'X':
+        direction = [direction[2], *direction[:2]]
+    elif axis == 'Y':
+        direction = [direction[0], -elevation, direction[1]]
+    direction = Vector(direction).normalized()
+    set_camera(direction, camera_dist=camera_dist)
+",point_e\evals\scripts\blender_script.py
+place_camera,,"def place_camera(time, camera_pose_mode='random', camera_dist_min=2.0,
+    camera_dist_max=2.0):
+    camera_dist = random.uniform(camera_dist_min, camera_dist_max)
+    if camera_pose_mode == 'random':
+        randomize_camera(camera_dist=camera_dist)
+    elif camera_pose_mode == 'z-circular':
+        pan_camera(time, axis='Z', camera_dist=camera_dist)
+    elif camera_pose_mode == 'z-circular-elevated':
+        pan_camera(time, axis='Z', camera_dist=camera_dist, elevation=
+            0.2617993878)
+    else:
+        raise ValueError(f'Unknown camera pose mode: {camera_pose_mode}')
+",point_e\evals\scripts\blender_script.py
+create_light,,"def create_light(location, energy=1.0, angle=0.5 * math.pi / 180):
+    light_data = bpy.data.lights.new(name='Light', type='SUN')
+    light_data.energy = energy
+    light_data.angle = angle
+    light_object = bpy.data.objects.new(name='Light', object_data=light_data)
+    direction = -location
+    rot_quat = direction.to_track_quat('-Z', 'Y')
+    light_object.rotation_euler = rot_quat.to_euler()
+    bpy.context.view_layer.update()
+    bpy.context.collection.objects.link(light_object)
+    light_object.location = location
+",point_e\evals\scripts\blender_script.py
+create_random_lights,,"def create_random_lights(count=4, distance=2.0, energy=1.5):
+    clear_lights()
+    for _ in range(count):
+        create_light(random_unit_vector() * distance, energy=energy)
+",point_e\evals\scripts\blender_script.py
+create_camera_light,,"def create_camera_light():
+    clear_lights()
+    create_light(bpy.context.scene.camera.location, energy=5.0)
+",point_e\evals\scripts\blender_script.py
+create_uniform_light,,"def create_uniform_light(backend):
+    clear_lights()
+    pos = Vector(UNIFORM_LIGHT_DIRECTION)
+    angle = 0.0092 if backend == 'CYCLES' else math.pi
+    create_light(pos, energy=5.0, angle=angle)
+    create_light(-pos, energy=5.0, angle=angle)
+",point_e\evals\scripts\blender_script.py
+create_vertex_color_shaders,,"def create_vertex_color_shaders():
+    for obj in bpy.context.scene.objects.values():
+        if not isinstance(obj.data, bpy.types.Mesh):
+            continue
+        if len(obj.data.materials):
+            continue
+        color_keys = (obj.data.vertex_colors or {}).keys()
+        if not len(color_keys):
+            continue
+        mat = bpy.data.materials.new(name='VertexColored')
+        mat.use_nodes = True
+        bsdf_node = None
+        for node in mat.node_tree.nodes:
+            if node.type == 'BSDF_PRINCIPLED':
+                bsdf_node = node
+        assert bsdf_node is not None, 'material has no Principled BSDF node to modify'
+        socket_map = {}
+        for input in bsdf_node.inputs:
+            socket_map[input.name] = input
+        socket_map['Specular'].default_value = 0.0
+        socket_map['Roughness'].default_value = 1.0
+        v_color = mat.node_tree.nodes.new('ShaderNodeVertexColor')
+        v_color.layer_name = color_keys[0]
+        mat.node_tree.links.new(v_color.outputs[0], socket_map['Base Color'])
+        obj.data.materials.append(mat)
+",point_e\evals\scripts\blender_script.py
+create_default_materials,,"def create_default_materials():
+    for obj in bpy.context.scene.objects.values():
+        if isinstance(obj.data, bpy.types.Mesh):
+            if not len(obj.data.materials):
+                mat = bpy.data.materials.new(name='DefaultMaterial')
+                mat.use_nodes = True
+                obj.data.materials.append(mat)
+",point_e\evals\scripts\blender_script.py
+find_materials,,"def find_materials():
+    all_materials = set()
+    for obj in bpy.context.scene.objects.values():
+        if not isinstance(obj.data, bpy.types.Mesh):
+            continue
+        for mat in obj.data.materials:
+            all_materials.add(mat)
+    return all_materials
+",point_e\evals\scripts\blender_script.py
+get_socket_value,,"def get_socket_value(tree, socket):
+    default = socket.default_value
+    if not isinstance(default, float):
+        default = list(default)
+    for link in tree.links:
+        if link.to_socket == socket:
+            return link.from_socket, default
+    return None, default
+",point_e\evals\scripts\blender_script.py
+clear_socket_input,,"def clear_socket_input(tree, socket):
+    for link in list(tree.links):
+        if link.to_socket == socket:
+            tree.links.remove(link)
+",point_e\evals\scripts\blender_script.py
+set_socket_value,,"def set_socket_value(tree, socket, socket_and_default):
+    clear_socket_input(tree, socket)
+    old_source_socket, default = socket_and_default
+    if isinstance(default, float) and not isinstance(socket.default_value,
+        float):
+        socket.default_value = [default] * 3 + [1.0]
+    else:
+        socket.default_value = default
+    if old_source_socket is not None:
+        tree.links.new(old_source_socket, socket)
+",point_e\evals\scripts\blender_script.py
+setup_nodes,,"def setup_nodes(output_path, capturing_material_alpha: bool=False):
+    tree = bpy.context.scene.node_tree
+    links = tree.links
+    for node in tree.nodes:
+        tree.nodes.remove(node)
+
+    def node_op(op: str, *args, clamp=False):
+        node = tree.nodes.new(type='CompositorNodeMath')
+        node.operation = op
+        if clamp:
+            node.use_clamp = True
+        for i, arg in enumerate(args):
+            if isinstance(arg, (int, float)):
+                node.inputs[i].default_value = arg
+            else:
+                links.new(arg, node.inputs[i])
+        return node.outputs[0]
+
+    def node_clamp(x, maximum=1.0):
+        return node_op('MINIMUM', x, maximum)
+
+    def node_mul(x, y, **kwargs):
+        return node_op('MULTIPLY', x, y, **kwargs)
+    input_node = tree.nodes.new(type='CompositorNodeRLayers')
+    input_node.scene = bpy.context.scene
+    input_sockets = {}
+    for output in input_node.outputs:
+        input_sockets[output.name] = output
+    if capturing_material_alpha:
+        color_socket = input_sockets['Image']
+    else:
+        raw_color_socket = input_sockets['Image']
+        color_node = tree.nodes.new(type='CompositorNodeConvertColorSpace')
+        color_node.from_color_space = 'Linear'
+        color_node.to_color_space = 'sRGB'
+        tree.links.new(raw_color_socket, color_node.inputs[0])
+        color_socket = color_node.outputs[0]
+    split_node = tree.nodes.new(type='CompositorNodeSepRGBA')
+    tree.links.new(color_socket, split_node.inputs[0])
+    for i, channel in (enumerate('rgba') if not capturing_material_alpha else
+        [(0, 'MatAlpha')]):
+        output_node = tree.nodes.new(type='CompositorNodeOutputFile')
+        output_node.base_path = f'{output_path}_{channel}'
+        links.new(split_node.outputs[i], output_node.inputs[0])
+    if capturing_material_alpha:
+        return
+    depth_out = node_clamp(node_mul(input_sockets['Depth'], 1 / MAX_DEPTH))
+    output_node = tree.nodes.new(type='CompositorNodeOutputFile')
+    output_node.base_path = f'{output_path}_depth'
+    links.new(depth_out, output_node.inputs[0])
+",point_e\evals\scripts\blender_script.py
+render_scene,,"def render_scene(output_path, fast_mode: bool):
+    use_workbench = bpy.context.scene.render.engine == 'BLENDER_WORKBENCH'
+    if use_workbench:
+        bpy.context.scene.render.engine = 'BLENDER_EEVEE'
+        bpy.context.scene.eevee.taa_render_samples = 1
+    if fast_mode:
+        if bpy.context.scene.render.engine == 'BLENDER_EEVEE':
+            bpy.context.scene.eevee.taa_render_samples = 1
+        elif bpy.context.scene.render.engine == 'CYCLES':
+            bpy.context.scene.cycles.samples = 256
+    elif bpy.context.scene.render.engine == 'CYCLES':
+        bpy.context.scene.cycles.time_limit = 40
+    bpy.context.view_layer.update()
+    bpy.context.scene.use_nodes = True
+    bpy.context.scene.view_layers['ViewLayer'].use_pass_z = True
+    bpy.context.scene.view_settings.view_transform = 'Raw'
+    bpy.context.scene.render.film_transparent = True
+    bpy.context.scene.render.resolution_x = 512
+    bpy.context.scene.render.resolution_y = 512
+    bpy.context.scene.render.image_settings.file_format = 'PNG'
+    bpy.context.scene.render.image_settings.color_mode = 'BW'
+    bpy.context.scene.render.image_settings.color_depth = '16'
+    bpy.context.scene.render.filepath = output_path
+    setup_nodes(output_path)
+    bpy.ops.render.render(write_still=True)
+    for channel_name in ['r', 'g', 'b', 'a', 'depth']:
+        sub_dir = f'{output_path}_{channel_name}'
+        image_path = os.path.join(sub_dir, os.listdir(sub_dir)[0])
+        name, ext = os.path.splitext(output_path)
+        if channel_name == 'depth' or not use_workbench:
+            os.rename(image_path, f'{name}_{channel_name}{ext}')
+        else:
+            os.remove(image_path)
+        os.removedirs(sub_dir)
+    if use_workbench:
+        bpy.context.scene.use_nodes = False
+        bpy.context.scene.render.engine = 'BLENDER_WORKBENCH'
+        bpy.context.scene.render.image_settings.color_mode = 'RGBA'
+        bpy.context.scene.render.image_settings.color_depth = '8'
+        bpy.context.scene.display.shading.color_type = 'TEXTURE'
+        bpy.context.scene.display.shading.light = 'FLAT'
+        if fast_mode:
+            bpy.context.scene.display.render_aa = 'FXAA'
+        os.remove(output_path)
+        bpy.ops.render.render(write_still=True)
+        bpy.context.scene.render.image_settings.color_mode = 'BW'
+        bpy.context.scene.render.image_settings.color_depth = '16'
+",point_e\evals\scripts\blender_script.py
+scene_fov,,"def scene_fov():
+    x_fov = bpy.context.scene.camera.data.angle_x
+    y_fov = bpy.context.scene.camera.data.angle_y
+    width = bpy.context.scene.render.resolution_x
+    height = bpy.context.scene.render.resolution_y
+    if bpy.context.scene.camera.data.angle == x_fov:
+        y_fov = 2 * math.atan(math.tan(x_fov / 2) * height / width)
+    else:
+        x_fov = 2 * math.atan(math.tan(y_fov / 2) * width / height)
+    return x_fov, y_fov
+",point_e\evals\scripts\blender_script.py
+write_camera_metadata,,"def write_camera_metadata(path):
+    x_fov, y_fov = scene_fov()
+    bbox_min, bbox_max = scene_bbox()
+    matrix = bpy.context.scene.camera.matrix_world
+    with open(path, 'w') as f:
+        json.dump(dict(format_version=FORMAT_VERSION, max_depth=MAX_DEPTH,
+            bbox=[list(bbox_min), list(bbox_max)], origin=list(matrix.col[3
+            ])[:3], x_fov=x_fov, y_fov=y_fov, x=list(matrix.col[0])[:3], y=
+            list(-matrix.col[1])[:3], z=list(-matrix.col[2])[:3]), f)
+",point_e\evals\scripts\blender_script.py
+save_rendering_dataset,,"def save_rendering_dataset(input_path: str, output_path: str, num_images:
+    int, backend: str, light_mode: str, camera_pose: str, camera_dist_min:
+    float, camera_dist_max: float, fast_mode: bool):
+    assert light_mode in ['random', 'uniform', 'camera']
+    assert camera_pose in ['random', 'z-circular', 'z-circular-elevated']
+    import_model(input_path)
+    bpy.context.scene.render.engine = backend
+    normalize_scene()
+    if light_mode == 'random':
+        create_random_lights()
+    elif light_mode == 'uniform':
+        create_uniform_light(backend)
+    create_camera()
+    create_vertex_color_shaders()
+    for i in range(num_images):
+        t = i / max(num_images - 1, 1)
+        place_camera(t, camera_pose_mode=camera_pose, camera_dist_min=
+            camera_dist_min, camera_dist_max=camera_dist_max)
+        if light_mode == 'camera':
+            create_camera_light()
+        render_scene(os.path.join(output_path, f'{i:05}.png'), fast_mode=
+            fast_mode)
+        write_camera_metadata(os.path.join(output_path, f'{i:05}.json'))
+    with open(os.path.join(output_path, 'info.json'), 'w') as f:
+        info = dict(backend=backend, light_mode=light_mode, fast_mode=
+            fast_mode, format_version=FORMAT_VERSION, channels=['R', 'G',
+            'B', 'A', 'D'], scale=0.5)
+        json.dump(info, f)
+",point_e\evals\scripts\blender_script.py
+main,,"def main():
+    try:
+        dash_index = sys.argv.index('--')
+    except ValueError as exc:
+        raise ValueError(""arguments must be preceded by '--'"") from exc
+    raw_args = sys.argv[dash_index + 1:]
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input_path', required=True, type=str)
+    parser.add_argument('--output_path', required=True, type=str)
+    parser.add_argument('--num_images', type=int, default=20)
+    parser.add_argument('--backend', type=str, default='BLENDER_EEVEE')
+    parser.add_argument('--light_mode', type=str, default='uniform')
+    parser.add_argument('--camera_pose', type=str, default='random')
+    parser.add_argument('--camera_dist_min', type=float, default=2.0)
+    parser.add_argument('--camera_dist_max', type=float, default=2.0)
+    parser.add_argument('--fast_mode', action='store_true')
+    args = parser.parse_args(raw_args)
+    save_rendering_dataset(input_path=args.input_path, output_path=args.
+        output_path, num_images=args.num_images, backend=args.backend,
+        light_mode=args.light_mode, camera_pose=args.camera_pose,
+        camera_dist_min=args.camera_dist_min, camera_dist_max=args.
+        camera_dist_max, fast_mode=args.fast_mode)
+",point_e\evals\scripts\blender_script.py
+main,,"def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--cache_dir', type=str, default=None)
+    parser.add_argument('batch_1', type=str)
+    parser.add_argument('batch_2', type=str)
+    args = parser.parse_args()
+    print('creating classifier...')
+    clf = PointNetClassifier(devices=get_torch_devices(), cache_dir=args.
+        cache_dir)
+    print('computing first batch activations')
+    features_1, _ = clf.features_and_preds(NpzStreamer(args.batch_1))
+    stats_1 = compute_statistics(features_1)
+    del features_1
+    features_2, _ = clf.features_and_preds(NpzStreamer(args.batch_2))
+    stats_2 = compute_statistics(features_2)
+    del features_2
+    print(f'P-FID: {stats_1.frechet_distance(stats_2)}')
+",point_e\evals\scripts\evaluate_pfid.py
+main,,"def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--cache_dir', type=str, default=None)
+    parser.add_argument('batch', type=str)
+    args = parser.parse_args()
+    print('creating classifier...')
+    clf = PointNetClassifier(devices=get_torch_devices(), cache_dir=args.
+        cache_dir)
+    print('computing batch predictions')
+    _, preds = clf.features_and_preds(NpzStreamer(args.batch))
+    print(f'P-IS: {compute_inception_score(preds)}')
+",point_e\evals\scripts\evaluate_pis.py
+checkpoint,"Evaluate a function without caching intermediate activations, allowing for
+reduced memory at the expense of extra compute in the backward pass.
+:param func: the function to evaluate.
+:param inputs: the argument sequence to pass to `func`.
+:param params: a sequence of parameters `func` depends on but does not
+               explicitly take as arguments.
+:param flag: if False, disable gradient checkpointing.","def checkpoint(func: Callable[..., Union[torch.Tensor, Sequence[torch.
+    Tensor]]], inputs: Sequence[torch.Tensor], params: Iterable[torch.
+    Tensor], flag: bool):
+    """"""""""""
+    if flag:
+        args = tuple(inputs) + tuple(params)
+        return CheckpointFunction.apply(func, len(inputs), *args)
+    else:
+        return func(*inputs)
+",point_e\models\checkpoint.py
+forward,,"@staticmethod
+def forward(ctx, run_function, length, *args):
+    ctx.run_function = run_function
+    ctx.input_tensors = list(args[:length])
+    ctx.input_params = list(args[length:])
+    with torch.no_grad():
+        output_tensors = ctx.run_function(*ctx.input_tensors)
+    return output_tensors
+",point_e\models\checkpoint.py
+backward,,"@staticmethod
+def backward(ctx, *output_grads):
+    ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.
+        input_tensors]
+    with torch.enable_grad():
+        shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
+        output_tensors = ctx.run_function(*shallow_copies)
+    input_grads = torch.autograd.grad(output_tensors, ctx.input_tensors +
+        ctx.input_params, output_grads, allow_unused=True)
+    del ctx.input_tensors
+    del ctx.input_params
+    del output_tensors
+    return (None, None) + input_grads
+",point_e\models\checkpoint.py
+model_from_config,,"def model_from_config(config: Dict[str, Any], device: torch.device
+    ) ->nn.Module:
+    config = config.copy()
+    name = config.pop('name')
+    if name == 'PointDiffusionTransformer':
+        return PointDiffusionTransformer(device=device, dtype=torch.float32,
+            **config)
+    elif name == 'CLIPImagePointDiffusionTransformer':
+        return CLIPImagePointDiffusionTransformer(device=device, dtype=
+            torch.float32, **config)
+    elif name == 'CLIPImageGridPointDiffusionTransformer':
+        return CLIPImageGridPointDiffusionTransformer(device=device, dtype=
+            torch.float32, **config)
+    elif name == 'UpsamplePointDiffusionTransformer':
+        return UpsamplePointDiffusionTransformer(device=device, dtype=torch
+            .float32, **config)
+    elif name == 'CLIPImageGridUpsamplePointDiffusionTransformer':
+        return CLIPImageGridUpsamplePointDiffusionTransformer(device=device,
+            dtype=torch.float32, **config)
+    elif name == 'CrossAttentionPointCloudSDFModel':
+        return CrossAttentionPointCloudSDFModel(device=device, dtype=torch.
+            float32, **config)
+    raise ValueError(f'unknown model name: {name}')
+",point_e\models\configs.py
+default_cache_dir,,"@lru_cache()
+def default_cache_dir() ->str:
+    return os.path.join(os.path.abspath(os.getcwd()), 'point_e_model_cache')
+",point_e\models\download.py
+fetch_file_cached,"Download the file at the given URL into a local file and return the path.
+If cache_dir is specified, it will be used to download the files.
+Otherwise, default_cache_dir() is used.","def fetch_file_cached(url: str, progress: bool=True, cache_dir: Optional[
+    str]=None, chunk_size: int=4096) ->str:
+    """"""""""""
+    if cache_dir is None:
+        cache_dir = default_cache_dir()
+    os.makedirs(cache_dir, exist_ok=True)
+    local_path = os.path.join(cache_dir, url.split('/')[-1])
+    if os.path.exists(local_path):
+        return local_path
+    response = requests.get(url, stream=True)
+    size = int(response.headers.get('content-length', '0'))
+    with FileLock(local_path + '.lock'):
+        if progress:
+            pbar = tqdm(total=size, unit='iB', unit_scale=True)
+        tmp_path = local_path + '.tmp'
+        with open(tmp_path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size):
+                if progress:
+                    pbar.update(len(chunk))
+                f.write(chunk)
+        os.rename(tmp_path, local_path)
+        if progress:
+            pbar.close()
+        return local_path
+",point_e\models\download.py
+load_checkpoint,,"def load_checkpoint(checkpoint_name: str, device: torch.device, progress:
+    bool=True, cache_dir: Optional[str]=None, chunk_size: int=4096) ->Dict[
+    str, torch.Tensor]:
+    if checkpoint_name not in MODEL_PATHS:
+        raise ValueError(
+            f'Unknown checkpoint name {checkpoint_name}. Known names are: {MODEL_PATHS.keys()}.'
+            )
+    path = fetch_file_cached(MODEL_PATHS[checkpoint_name], progress=
+        progress, cache_dir=cache_dir, chunk_size=chunk_size)
+    return torch.load(path, map_location=device)
+",point_e\models\download.py
+__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_data: int,
+    width: int, heads: int, init_scale: float, data_width: Optional[int]=None):
+    super().__init__()
+    self.n_data = n_data
+    self.width = width
+    self.heads = heads
+    self.data_width = width if data_width is None else data_width
+    self.c_q = nn.Linear(width, width, device=device, dtype=dtype)
+    self.c_kv = nn.Linear(self.data_width, width * 2, device=device, dtype=
+        dtype)
+    self.c_proj = nn.Linear(width, width, device=device, dtype=dtype)
+    self.attention = QKVMultiheadCrossAttention(device=device, dtype=dtype,
+        heads=heads, n_data=n_data)
+    init_linear(self.c_q, init_scale)
+    init_linear(self.c_kv, init_scale)
+    init_linear(self.c_proj, init_scale)
+",point_e\models\perceiver.py
+forward,,"def forward(self, x, data):
+    x = self.c_q(x)
+    data = self.c_kv(data)
+    x = checkpoint(self.attention, (x, data), (), True)
+    x = self.c_proj(x)
+    return x
+",point_e\models\perceiver.py
+__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, heads: int,
+    n_data: int):
+    super().__init__()
+    self.device = device
+    self.dtype = dtype
+    self.heads = heads
+    self.n_data = n_data
+",point_e\models\perceiver.py
+forward,,"def forward(self, q, kv):
+    _, n_ctx, _ = q.shape
+    bs, n_data, width = kv.shape
+    attn_ch = width // self.heads // 2
+    scale = 1 / math.sqrt(math.sqrt(attn_ch))
+    q = q.view(bs, n_ctx, self.heads, -1)
+    kv = kv.view(bs, n_data, self.heads, -1)
+    k, v = torch.split(kv, attn_ch, dim=-1)
+    weight = torch.einsum('bthc,bshc->bhts', q * scale, k * scale)
+    wdtype = weight.dtype
+    weight = torch.softmax(weight.float(), dim=-1).type(wdtype)
+    return torch.einsum('bhts,bshc->bthc', weight, v).reshape(bs, n_ctx, -1)
+",point_e\models\perceiver.py
+__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_data: int,
+    width: int, heads: int, data_width: Optional[int]=None, init_scale:
+    float=1.0):
+    super().__init__()
+    if data_width is None:
+        data_width = width
+    self.attn = MultiheadCrossAttention(device=device, dtype=dtype, n_data=
+        n_data, width=width, heads=heads, data_width=data_width, init_scale
+        =init_scale)
+    self.ln_1 = nn.LayerNorm(width, device=device, dtype=dtype)
+    self.ln_2 = nn.LayerNorm(data_width, device=device, dtype=dtype)
+    self.mlp = MLP(device=device, dtype=dtype, width=width, init_scale=
+        init_scale)
+    self.ln_3 = nn.LayerNorm(width, device=device, dtype=dtype)
+",point_e\models\perceiver.py
+forward,,"def forward(self, x: torch.Tensor, data: torch.Tensor):
+    x = x + self.attn(self.ln_1(x), self.ln_2(data))
+    x = x + self.mlp(self.ln_3(x))
+    return x
+",point_e\models\perceiver.py
+__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_data: int,
+    width: int, layers: int, heads: int, init_scale: float=0.25, data_width:
+    Optional[int]=None):
+    super().__init__()
+    self.width = width
+    self.layers = layers
+    init_scale = init_scale * math.sqrt(1.0 / width)
+    self.resblocks = nn.ModuleList([ResidualCrossAttentionBlock(device=
+        device, dtype=dtype, n_data=n_data, width=width, heads=heads,
+        init_scale=init_scale, data_width=data_width) for _ in range(layers)])
+",point_e\models\perceiver.py
+forward,,"def forward(self, x: torch.Tensor, data: torch.Tensor):
+    for block in self.resblocks:
+        x = block(x, data)
+    return x
+",point_e\models\perceiver.py
+_image_to_pil,,"def _image_to_pil(obj: Optional[ImageType]) ->Image.Image:
+    if obj is None:
+        return Image.fromarray(np.zeros([64, 64, 3], dtype=np.uint8))
+    if isinstance(obj, np.ndarray):
+        return Image.fromarray(obj.astype(np.uint8))
+    elif isinstance(obj, torch.Tensor):
+        return Image.fromarray(obj.detach().cpu().numpy().astype(np.uint8))
+    else:
+        return obj
+",point_e\models\pretrained_clip.py
+__init__,,"def __init__(self, device: torch.device, dtype: Optional[torch.dtype]=torch
+    .float32, ensure_used_params: bool=True, clip_name: str='ViT-L/14',
+    cache_dir: Optional[str]=None):
+    super().__init__()
+    assert clip_name in ['ViT-L/14', 'ViT-B/32']
+    self.device = device
+    self.ensure_used_params = ensure_used_params
+    import clip
+    self.clip_model, self.preprocess = clip.load(clip_name, device=device,
+        download_root=cache_dir or default_cache_dir())
+    self.clip_name = clip_name
+    if dtype is not None:
+        self.clip_model.to(dtype)
+    self._tokenize = clip.tokenize
+",point_e\models\pretrained_clip.py
+feature_dim,,"@property
+def feature_dim(self) ->int:
+    if self.clip_name == 'ViT-L/14':
+        return 768
+    else:
+        return 512
+",point_e\models\pretrained_clip.py
+grid_size,,"@property
+def grid_size(self) ->int:
+    if self.clip_name == 'ViT-L/14':
+        return 16
+    else:
+        return 7
+",point_e\models\pretrained_clip.py
+grid_feature_dim,,"@property
+def grid_feature_dim(self) ->int:
+    if self.clip_name == 'ViT-L/14':
+        return 1024
+    else:
+        return 768
+",point_e\models\pretrained_clip.py
+forward,"Generate a batch of embeddings from a mixture of images, texts,
+precomputed embeddings, and possibly empty values.
+
+For each batch element, at most one of images, texts, and embeddings
+should have a non-None value. Embeddings from multiple modalities
+cannot be mixed for a single batch element. If no modality is provided,
+a zero embedding will be used for the batch element.","def forward(self, batch_size: int, images: Optional[Iterable[Optional[
+    ImageType]]]=None, texts: Optional[Iterable[Optional[str]]]=None,
+    embeddings: Optional[Iterable[Optional[torch.Tensor]]]=None
+    ) ->torch.Tensor:
+    """"""""""""
+    image_seq = [None] * batch_size if images is None else list(images)
+    text_seq = [None] * batch_size if texts is None else list(texts)
+    embedding_seq = [None] * batch_size if embeddings is None else list(
+        embeddings)
+    assert len(image_seq
+        ) == batch_size, 'number of images should match batch size'
+    assert len(text_seq
+        ) == batch_size, 'number of texts should match batch size'
+    assert len(embedding_seq
+        ) == batch_size, 'number of embeddings should match batch size'
+    if self.ensure_used_params:
+        return self._static_multimodal_embed(images=image_seq, texts=
+            text_seq, embeddings=embedding_seq)
+    result = torch.zeros((batch_size, self.feature_dim), device=self.device)
+    index_images = []
+    index_texts = []
+    for i, (image, text, emb) in enumerate(zip(image_seq, text_seq,
+        embedding_seq)):
+        assert sum([int(image is not None), int(text is not None), int(emb
+             is not None)]
+            ) < 2, 'only one modality may be non-None per batch element'
+        if image is not None:
+            index_images.append((i, image))
+        elif text is not None:
+            index_texts.append((i, text))
+        elif emb is not None:
+            result[i] = emb.to(result)
+    if len(index_images):
+        embs = self.embed_images(img for _, img in index_images)
+        for (i, _), emb in zip(index_images, embs):
+            result[i] = emb.to(result)
+    if len(index_texts):
+        embs = self.embed_text(text for _, text in index_texts)
+        for (i, _), emb in zip(index_texts, embs):
+            result[i] = emb.to(result)
+    return result
+",point_e\models\pretrained_clip.py
+_static_multimodal_embed,"Like forward(), but always runs all encoders to ensure that
+the forward graph looks the same on every rank.","def _static_multimodal_embed(self, images: List[Optional[ImageType]]=None,
+    texts: List[Optional[str]]=None, embeddings: List[Optional[torch.Tensor
+    ]]=None) ->torch.Tensor:
+    """"""""""""
+    image_emb = self.embed_images(images)
+    text_emb = self.embed_text(t if t else '' for t in texts)
+    joined_embs = torch.stack([(emb.to(device=self.device, dtype=torch.
+        float32) if emb is not None else torch.zeros(self.feature_dim,
+        device=self.device)) for emb in embeddings], dim=0)
+    image_flag = torch.tensor([(x is not None) for x in images], device=
+        self.device)[:, None].expand_as(image_emb)
+    text_flag = torch.tensor([(x is not None) for x in texts], device=self.
+        device)[:, None].expand_as(image_emb)
+    emb_flag = torch.tensor([(x is not None) for x in embeddings], device=
+        self.device)[:, None].expand_as(image_emb)
+    return image_flag.float() * image_emb + text_flag.float(
+        ) * text_emb + emb_flag.float(
+        ) * joined_embs + self.clip_model.logit_scale * 0
+",point_e\models\pretrained_clip.py
+embed_images,":param xs: N images, stored as numpy arrays, tensors, or PIL images.
+:return: an [N x D] tensor of features.","def embed_images(self, xs: Iterable[Optional[ImageType]]) ->torch.Tensor:
+    """"""""""""
+    clip_inputs = self.images_to_tensor(xs)
+    results = self.clip_model.encode_image(clip_inputs).float()
+    return results / torch.linalg.norm(results, dim=-1, keepdim=True)
+",point_e\models\pretrained_clip.py
+embed_text,Embed text prompts as an [N x D] tensor.,"def embed_text(self, prompts: Iterable[str]) ->torch.Tensor:
+    """"""""""""
+    enc = self.clip_model.encode_text(self._tokenize(list(prompts),
+        truncate=True).to(self.device)).float()
+    return enc / torch.linalg.norm(enc, dim=-1, keepdim=True)
+",point_e\models\pretrained_clip.py
+embed_images_grid,"Embed images into latent grids.
+
+:param xs: an iterable of images to embed.
+:return: a tensor of shape [N x C x L], where L = self.grid_size**2.","def embed_images_grid(self, xs: Iterable[Optional[ImageType]]) ->torch.Tensor:
+    """"""""""""
+    if self.ensure_used_params:
+        extra_value = 0.0
+        for p in self.parameters():
+            extra_value = extra_value + p.mean() * 0.0
+    else:
+        extra_value = 0.0
+    x = self.images_to_tensor(xs).to(self.clip_model.dtype)
+    vt = self.clip_model.visual
+    x = vt.conv1(x)
+    x = x.reshape(x.shape[0], x.shape[1], -1)
+    x = x.permute(0, 2, 1)
+    x = torch.cat([vt.class_embedding.to(x.dtype) + torch.zeros(x.shape[0],
+        1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)
+    x = x + vt.positional_embedding.to(x.dtype)
+    x = vt.ln_pre(x)
+    x = x.permute(1, 0, 2)
+    x = vt.transformer(x)
+    x = x.permute(1, 2, 0)
+    return x[..., 1:].contiguous().float() + extra_value
+",point_e\models\pretrained_clip.py
+images_to_tensor,,"def images_to_tensor(self, xs: Iterable[Optional[ImageType]]) ->torch.Tensor:
+    return torch.stack([self.preprocess(_image_to_pil(x)) for x in xs], dim=0
+        ).to(self.device)
+",point_e\models\pretrained_clip.py
+__init__,,"def __init__(self, device: torch.device, **kwargs):
+    self.model = ImageCLIP(device, dtype=None, ensure_used_params=False, **
+        kwargs)
+    for parameter in self.model.parameters():
+        parameter.requires_grad_(False)
+",point_e\models\pretrained_clip.py
+feature_dim,,"@property
+def feature_dim(self) ->int:
+    return self.model.feature_dim
+",point_e\models\pretrained_clip.py
+grid_size,,"@property
+def grid_size(self) ->int:
+    return self.model.grid_size
+",point_e\models\pretrained_clip.py
+grid_feature_dim,,"@property
+def grid_feature_dim(self) ->int:
+    return self.model.grid_feature_dim
+",point_e\models\pretrained_clip.py
+__call__,,"def __call__(self, batch_size: int, images: Optional[Iterable[Optional[
+    ImageType]]]=None, texts: Optional[Iterable[Optional[str]]]=None,
+    embeddings: Optional[Iterable[Optional[torch.Tensor]]]=None
+    ) ->torch.Tensor:
+    return self.model(batch_size=batch_size, images=images, texts=texts,
+        embeddings=embeddings)
+",point_e\models\pretrained_clip.py
+embed_images,,"def embed_images(self, xs: Iterable[Optional[ImageType]]) ->torch.Tensor:
+    with torch.no_grad():
+        return self.model.embed_images(xs)
+",point_e\models\pretrained_clip.py
+embed_text,,"def embed_text(self, prompts: Iterable[str]) ->torch.Tensor:
+    with torch.no_grad():
+        return self.model.embed_text(prompts)
+",point_e\models\pretrained_clip.py
+embed_images_grid,,"def embed_images_grid(self, xs: Iterable[Optional[ImageType]]) ->torch.Tensor:
+    with torch.no_grad():
+        return self.model.embed_images_grid(xs)
+",point_e\models\pretrained_clip.py
+device,Get the device that should be used for input tensors.,"@property
+@abstractmethod
+def device(self) ->torch.device:
+    """"""""""""
+",point_e\models\sdf.py
+default_batch_size,"Get a reasonable default number of query points for the model.
+In some cases, this might be the only supported size.","@property
+@abstractmethod
+def default_batch_size(self) ->int:
+    """"""""""""
+",point_e\models\sdf.py
+encode_point_clouds,"Encode a batch of point clouds to cache part of the SDF calculation
+done by forward().
+
+:param point_clouds: a batch of [batch x 3 x N] points.
+:return: a state representing the encoded point cloud batch.","@abstractmethod
+def encode_point_clouds(self, point_clouds: torch.Tensor) ->Dict[str, torch
+    .Tensor]:
+    """"""""""""
+",point_e\models\sdf.py
+forward,"Predict the SDF at the coordinates x, given a batch of point clouds.
+
+Either point_clouds or encoded should be passed. Only exactly one of
+these arguments should be None.
+
+:param x: a [batch x 3 x N'] tensor of query points.
+:param point_clouds: a [batch x 3 x N] batch of point clouds.
+:param encoded: the result of calling encode_point_clouds().
+:return: a [batch x N'] tensor of SDF predictions.","def forward(self, x: torch.Tensor, point_clouds: Optional[torch.Tensor]=
+    None, encoded: Optional[Dict[str, torch.Tensor]]=None) ->torch.Tensor:
+    """"""""""""
+    assert point_clouds is not None or encoded is not None
+    assert point_clouds is None or encoded is None
+    if point_clouds is not None:
+        encoded = self.encode_point_clouds(point_clouds)
+    return self.predict_sdf(x, encoded)
+",point_e\models\sdf.py
+predict_sdf,"Predict the SDF at the query points given the encoded point clouds.
+
+Each query point should be treated independently, only conditioning on
+the point clouds themselves.","@abstractmethod
+def predict_sdf(self, x: torch.Tensor, encoded: Optional[Dict[str, torch.
+    Tensor]]) ->torch.Tensor:
+    """"""""""""
+",point_e\models\sdf.py
+__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int=
+    4096, width: int=512, encoder_layers: int=12, encoder_heads: int=8,
+    decoder_layers: int=4, decoder_heads: int=8, init_scale: float=0.25):
+    super().__init__()
+    self._device = device
+    self.n_ctx = n_ctx
+    self.encoder_input_proj = nn.Linear(3, width, device=device, dtype=dtype)
+    self.encoder = Transformer(device=device, dtype=dtype, n_ctx=n_ctx,
+        width=width, layers=encoder_layers, heads=encoder_heads, init_scale
+        =init_scale)
+    self.decoder_input_proj = nn.Linear(3, width, device=device, dtype=dtype)
+    self.decoder = SimplePerceiver(device=device, dtype=dtype, n_data=n_ctx,
+        width=width, layers=decoder_layers, heads=decoder_heads, init_scale
+        =init_scale)
+    self.ln_post = nn.LayerNorm(width, device=device, dtype=dtype)
+    self.output_proj = nn.Linear(width, 1, device=device, dtype=dtype)
+",point_e\models\sdf.py
+device,,"@property
+def device(self) ->torch.device:
+    return self._device
+",point_e\models\sdf.py
+default_batch_size,,"@property
+def default_batch_size(self) ->int:
+    return self.n_query
+",point_e\models\sdf.py
+encode_point_clouds,,"def encode_point_clouds(self, point_clouds: torch.Tensor) ->Dict[str, torch
+    .Tensor]:
+    h = self.encoder_input_proj(point_clouds.permute(0, 2, 1))
+    h = self.encoder(h)
+    return dict(latents=h)
+",point_e\models\sdf.py
+predict_sdf,,"def predict_sdf(self, x: torch.Tensor, encoded: Optional[Dict[str, torch.
+    Tensor]]) ->torch.Tensor:
+    data = encoded['latents']
+    x = self.decoder_input_proj(x.permute(0, 2, 1))
+    x = self.decoder(x, data)
+    x = self.ln_post(x)
+    x = self.output_proj(x)
+    return x[..., 0]
+",point_e\models\sdf.py
+init_linear,,"def init_linear(l, stddev):
+    nn.init.normal_(l.weight, std=stddev)
+    if l.bias is not None:
+        nn.init.constant_(l.bias, 0.0)
+",point_e\models\transformer.py
+__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int,
+    width: int, heads: int, init_scale: float):
+    super().__init__()
+    self.n_ctx = n_ctx
+    self.width = width
+    self.heads = heads
+    self.c_qkv = nn.Linear(width, width * 3, device=device, dtype=dtype)
+    self.c_proj = nn.Linear(width, width, device=device, dtype=dtype)
+    self.attention = QKVMultiheadAttention(device=device, dtype=dtype,
+        heads=heads, n_ctx=n_ctx)
+    init_linear(self.c_qkv, init_scale)
+    init_linear(self.c_proj, init_scale)
+",point_e\models\transformer.py
+forward,,"def forward(self, x):
+    x = self.c_qkv(x)
+    x = checkpoint(self.attention, (x,), (), True)
+    x = self.c_proj(x)
+    return x
+",point_e\models\transformer.py
+__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, width: int,
+    init_scale: float):
+    super().__init__()
+    self.width = width
+    self.c_fc = nn.Linear(width, width * 4, device=device, dtype=dtype)
+    self.c_proj = nn.Linear(width * 4, width, device=device, dtype=dtype)
+    self.gelu = nn.GELU()
+    init_linear(self.c_fc, init_scale)
+    init_linear(self.c_proj, init_scale)
+",point_e\models\transformer.py
+forward,,"def forward(self, x):
+    return self.c_proj(self.gelu(self.c_fc(x)))
+",point_e\models\transformer.py
+__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, heads: int,
+    n_ctx: int):
+    super().__init__()
+    self.device = device
+    self.dtype = dtype
+    self.heads = heads
+    self.n_ctx = n_ctx
+",point_e\models\transformer.py
+forward,,"def forward(self, qkv):
+    bs, n_ctx, width = qkv.shape
+    attn_ch = width // self.heads // 3
+    scale = 1 / math.sqrt(math.sqrt(attn_ch))
+    qkv = qkv.view(bs, n_ctx, self.heads, -1)
+    q, k, v = torch.split(qkv, attn_ch, dim=-1)
+    weight = torch.einsum('bthc,bshc->bhts', q * scale, k * scale)
+    wdtype = weight.dtype
+    weight = torch.softmax(weight.float(), dim=-1).type(wdtype)
+    return torch.einsum('bhts,bshc->bthc', weight, v).reshape(bs, n_ctx, -1)
+",point_e\models\transformer.py
+__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int,
+    width: int, heads: int, init_scale: float=1.0):
+    super().__init__()
+    self.attn = MultiheadAttention(device=device, dtype=dtype, n_ctx=n_ctx,
+        width=width, heads=heads, init_scale=init_scale)
+    self.ln_1 = nn.LayerNorm(width, device=device, dtype=dtype)
+    self.mlp = MLP(device=device, dtype=dtype, width=width, init_scale=
+        init_scale)
+    self.ln_2 = nn.LayerNorm(width, device=device, dtype=dtype)
+",point_e\models\transformer.py
+forward,,"def forward(self, x: torch.Tensor):
+    x = x + self.attn(self.ln_1(x))
+    x = x + self.mlp(self.ln_2(x))
+    return x
+",point_e\models\transformer.py
+__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int,
+    width: int, layers: int, heads: int, init_scale: float=0.25):
+    super().__init__()
+    self.n_ctx = n_ctx
+    self.width = width
+    self.layers = layers
+    init_scale = init_scale * math.sqrt(1.0 / width)
+    self.resblocks = nn.ModuleList([ResidualAttentionBlock(device=device,
+        dtype=dtype, n_ctx=n_ctx, width=width, heads=heads, init_scale=
+        init_scale) for _ in range(layers)])
+",point_e\models\transformer.py
+forward,,"def forward(self, x: torch.Tensor):
+    for block in self.resblocks:
+        x = block(x)
+    return x
+",point_e\models\transformer.py
+__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype,
+    input_channels: int=3, output_channels: int=3, n_ctx: int=1024, width:
+    int=512, layers: int=12, heads: int=8, init_scale: float=0.25,
+    time_token_cond: bool=False):
+    super().__init__()
+    self.input_channels = input_channels
+    self.output_channels = output_channels
+    self.n_ctx = n_ctx
+    self.time_token_cond = time_token_cond
+    self.time_embed = MLP(device=device, dtype=dtype, width=width,
+        init_scale=init_scale * math.sqrt(1.0 / width))
+    self.ln_pre = nn.LayerNorm(width, device=device, dtype=dtype)
+    self.backbone = Transformer(device=device, dtype=dtype, n_ctx=n_ctx +
+        int(time_token_cond), width=width, layers=layers, heads=heads,
+        init_scale=init_scale)
+    self.ln_post = nn.LayerNorm(width, device=device, dtype=dtype)
+    self.input_proj = nn.Linear(input_channels, width, device=device, dtype
+        =dtype)
+    self.output_proj = nn.Linear(width, output_channels, device=device,
+        dtype=dtype)
+    with torch.no_grad():
+        self.output_proj.weight.zero_()
+        self.output_proj.bias.zero_()
+",point_e\models\transformer.py
+forward,":param x: an [N x C x T] tensor.
+:param t: an [N] tensor.
+:return: an [N x C' x T] tensor.","def forward(self, x: torch.Tensor, t: torch.Tensor):
+    """"""""""""
+    assert x.shape[-1] == self.n_ctx
+    t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
+    return self._forward_with_cond(x, [(t_embed, self.time_token_cond)])
+",point_e\models\transformer.py
+_forward_with_cond,,"def _forward_with_cond(self, x: torch.Tensor, cond_as_token: List[Tuple[
+    torch.Tensor, bool]]) ->torch.Tensor:
+    h = self.input_proj(x.permute(0, 2, 1))
+    for emb, as_token in cond_as_token:
+        if not as_token:
+            h = h + emb[:, None]
+    extra_tokens = [(emb[:, None] if len(emb.shape) == 2 else emb) for emb,
+        as_token in cond_as_token if as_token]
+    if len(extra_tokens):
+        h = torch.cat(extra_tokens + [h], dim=1)
+    h = self.ln_pre(h)
+    h = self.backbone(h)
+    h = self.ln_post(h)
+    if len(extra_tokens):
+        h = h[:, sum(h.shape[1] for h in extra_tokens):]
+    h = self.output_proj(h)
+    return h.permute(0, 2, 1)
+",point_e\models\transformer.py
+__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int=
+    1024, token_cond: bool=False, cond_drop_prob: float=0.0, frozen_clip:
+    bool=True, cache_dir: Optional[str]=None, **kwargs):
+    super().__init__(device=device, dtype=dtype, n_ctx=n_ctx + int(
+        token_cond), **kwargs)
+    self.n_ctx = n_ctx
+    self.token_cond = token_cond
+    self.clip = (FrozenImageCLIP if frozen_clip else ImageCLIP)(device,
+        cache_dir=cache_dir)
+    self.clip_embed = nn.Linear(self.clip.feature_dim, self.backbone.width,
+        device=device, dtype=dtype)
+    self.cond_drop_prob = cond_drop_prob
+",point_e\models\transformer.py
+cached_model_kwargs,,"def cached_model_kwargs(self, batch_size: int, model_kwargs: Dict[str, Any]
+    ) ->Dict[str, Any]:
+    with torch.no_grad():
+        return dict(embeddings=self.clip(batch_size, **model_kwargs))
+",point_e\models\transformer.py
+forward,":param x: an [N x C x T] tensor.
+:param t: an [N] tensor.
+:param images: a batch of images to condition on.
+:param texts: a batch of texts to condition on.
+:param embeddings: a batch of CLIP embeddings to condition on.
+:return: an [N x C' x T] tensor.","def forward(self, x: torch.Tensor, t: torch.Tensor, images: Optional[
+    Iterable[Optional[ImageType]]]=None, texts: Optional[Iterable[Optional[
+    str]]]=None, embeddings: Optional[Iterable[Optional[torch.Tensor]]]=None):
+    """"""""""""
+    assert x.shape[-1] == self.n_ctx
+    t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
+    clip_out = self.clip(batch_size=len(x), images=images, texts=texts,
+        embeddings=embeddings)
+    assert len(clip_out.shape) == 2 and clip_out.shape[0] == x.shape[0]
+    if self.training:
+        mask = torch.rand(size=[len(x)]) >= self.cond_drop_prob
+        clip_out = clip_out * mask[:, None].to(clip_out)
+    clip_out = math.sqrt(clip_out.shape[1]) * clip_out
+    clip_embed = self.clip_embed(clip_out)
+    cond = [(clip_embed, self.token_cond), (t_embed, self.time_token_cond)]
+    return self._forward_with_cond(x, cond)
+",point_e\models\transformer.py
+__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int=
+    1024, cond_drop_prob: float=0.0, frozen_clip: bool=True, cache_dir:
+    Optional[str]=None, **kwargs):
+    clip = (FrozenImageCLIP if frozen_clip else ImageCLIP)(device,
+        cache_dir=cache_dir)
+    super().__init__(device=device, dtype=dtype, n_ctx=n_ctx + clip.
+        grid_size ** 2, **kwargs)
+    self.n_ctx = n_ctx
+    self.clip = clip
+    self.clip_embed = nn.Sequential(nn.LayerNorm(normalized_shape=(self.
+        clip.grid_feature_dim,), device=device, dtype=dtype), nn.Linear(
+        self.clip.grid_feature_dim, self.backbone.width, device=device,
+        dtype=dtype))
+    self.cond_drop_prob = cond_drop_prob
+",point_e\models\transformer.py
+cached_model_kwargs,,"def cached_model_kwargs(self, batch_size: int, model_kwargs: Dict[str, Any]
+    ) ->Dict[str, Any]:
+    _ = batch_size
+    with torch.no_grad():
+        return dict(embeddings=self.clip.embed_images_grid(model_kwargs[
+            'images']))
+",point_e\models\transformer.py
+forward,":param x: an [N x C x T] tensor.
+:param t: an [N] tensor.
+:param images: a batch of images to condition on.
+:param embeddings: a batch of CLIP latent grids to condition on.
+:return: an [N x C' x T] tensor.","def forward(self, x: torch.Tensor, t: torch.Tensor, images: Optional[
+    Iterable[ImageType]]=None, embeddings: Optional[Iterable[torch.Tensor]]
+    =None):
+    """"""""""""
+    assert images is not None or embeddings is not None, 'must specify images or embeddings'
+    assert images is None or embeddings is None, 'cannot specify both images and embeddings'
+    assert x.shape[-1] == self.n_ctx
+    t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
+    if images is not None:
+        clip_out = self.clip.embed_images_grid(images)
+    else:
+        clip_out = embeddings
+    if self.training:
+        mask = torch.rand(size=[len(x)]) >= self.cond_drop_prob
+        clip_out = clip_out * mask[:, None, None].to(clip_out)
+    clip_out = clip_out.permute(0, 2, 1)
+    clip_embed = self.clip_embed(clip_out)
+    cond = [(t_embed, self.time_token_cond), (clip_embed, True)]
+    return self._forward_with_cond(x, cond)
+",point_e\models\transformer.py
+__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype,
+    cond_input_channels: Optional[int]=None, cond_ctx: int=1024, n_ctx: int
+    =4096 - 1024, channel_scales: Optional[Sequence[float]]=None,
+    channel_biases: Optional[Sequence[float]]=None, **kwargs):
+    super().__init__(device=device, dtype=dtype, n_ctx=n_ctx + cond_ctx, **
+        kwargs)
+    self.n_ctx = n_ctx
+    self.cond_input_channels = cond_input_channels or self.input_channels
+    self.cond_point_proj = nn.Linear(self.cond_input_channels, self.
+        backbone.width, device=device, dtype=dtype)
+    self.register_buffer('channel_scales', torch.tensor(channel_scales,
+        dtype=dtype, device=device) if channel_scales is not None else None)
+    self.register_buffer('channel_biases', torch.tensor(channel_biases,
+        dtype=dtype, device=device) if channel_biases is not None else None)
+",point_e\models\transformer.py
+forward,":param x: an [N x C1 x T] tensor.
+:param t: an [N] tensor.
+:param low_res: an [N x C2 x T'] tensor of conditioning points.
+:return: an [N x C3 x T] tensor.","def forward(self, x: torch.Tensor, t: torch.Tensor, *, low_res: torch.Tensor):
+    """"""""""""
+    assert x.shape[-1] == self.n_ctx
+    t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
+    low_res_embed = self._embed_low_res(low_res)
+    cond = [(t_embed, self.time_token_cond), (low_res_embed, True)]
+    return self._forward_with_cond(x, cond)
+",point_e\models\transformer.py
+_embed_low_res,,"def _embed_low_res(self, x: torch.Tensor) ->torch.Tensor:
+    if self.channel_scales is not None:
+        x = x * self.channel_scales[None, :, None]
+    if self.channel_biases is not None:
+        x = x + self.channel_biases[None, :, None]
+    return self.cond_point_proj(x.permute(0, 2, 1))
+",point_e\models\transformer.py
+__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int=
+    4096 - 1024, cond_drop_prob: float=0.0, frozen_clip: bool=True,
+    cache_dir: Optional[str]=None, **kwargs):
+    clip = (FrozenImageCLIP if frozen_clip else ImageCLIP)(device,
+        cache_dir=cache_dir)
+    super().__init__(device=device, dtype=dtype, n_ctx=n_ctx + clip.
+        grid_size ** 2, **kwargs)
+    self.n_ctx = n_ctx
+    self.clip = clip
+    self.clip_embed = nn.Sequential(nn.LayerNorm(normalized_shape=(self.
+        clip.grid_feature_dim,), device=device, dtype=dtype), nn.Linear(
+        self.clip.grid_feature_dim, self.backbone.width, device=device,
+        dtype=dtype))
+    self.cond_drop_prob = cond_drop_prob
+",point_e\models\transformer.py
+cached_model_kwargs,,"def cached_model_kwargs(self, batch_size: int, model_kwargs: Dict[str, Any]
+    ) ->Dict[str, Any]:
+    if 'images' not in model_kwargs:
+        zero_emb = torch.zeros([batch_size, self.clip.grid_feature_dim, 
+            self.clip.grid_size ** 2], device=next(self.parameters()).device)
+        return dict(embeddings=zero_emb, low_res=model_kwargs['low_res'])
+    with torch.no_grad():
+        return dict(embeddings=self.clip.embed_images_grid(model_kwargs[
+            'images']), low_res=model_kwargs['low_res'])
+",point_e\models\transformer.py
+forward,":param x: an [N x C1 x T] tensor.
+:param t: an [N] tensor.
+:param low_res: an [N x C2 x T'] tensor of conditioning points.
+:param images: a batch of images to condition on.
+:param embeddings: a batch of CLIP latent grids to condition on.
+:return: an [N x C3 x T] tensor.","def forward(self, x: torch.Tensor, t: torch.Tensor, *, low_res: torch.
+    Tensor, images: Optional[Iterable[ImageType]]=None, embeddings:
+    Optional[Iterable[torch.Tensor]]=None):
+    """"""""""""
+    assert x.shape[-1] == self.n_ctx
+    t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
+    low_res_embed = self._embed_low_res(low_res)
+    if images is not None:
+        clip_out = self.clip.embed_images_grid(images)
+    elif embeddings is not None:
+        clip_out = embeddings
+    else:
+        clip_out = torch.zeros([len(x), self.clip.grid_feature_dim, self.
+            clip.grid_size ** 2], dtype=x.dtype, device=x.device)
+    if self.training:
+        mask = torch.rand(size=[len(x)]) >= self.cond_drop_prob
+        clip_out = clip_out * mask[:, None, None].to(clip_out)
+    clip_out = clip_out.permute(0, 2, 1)
+    clip_embed = self.clip_embed(clip_out)
+    cond = [(t_embed, self.time_token_cond), (clip_embed, True), (
+        low_res_embed, True)]
+    return self._forward_with_cond(x, cond)
+",point_e\models\transformer.py
+timestep_embedding,"Create sinusoidal timestep embeddings.
+:param timesteps: a 1-D Tensor of N indices, one per batch element.
+                  These may be fractional.
+:param dim: the dimension of the output.
+:param max_period: controls the minimum frequency of the embeddings.
+:return: an [N x dim] Tensor of positional embeddings.","def timestep_embedding(timesteps, dim, max_period=10000):
+    """"""""""""
+    half = dim // 2
+    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=
+        half, dtype=torch.float32) / half).to(device=timesteps.device)
+    args = timesteps[:, None].to(timesteps.dtype) * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1]
+            )], dim=-1)
+    return embedding
+",point_e\models\util.py
+load,Load the mesh from a .npz file.,"@classmethod
+def load(cls, f: Union[str, BinaryIO]) ->'TriMesh':
+    """"""""""""
+    if isinstance(f, str):
+        with open(f, 'rb') as reader:
+            return cls.load(reader)
+    else:
+        obj = np.load(f)
+        keys = list(obj.keys())
+        verts = obj['verts']
+        faces = obj['faces']
+        normals = obj['normals'] if 'normals' in keys else None
+        vertex_channels = {}
+        face_channels = {}
+        for key in keys:
+            if key.startswith('v_'):
+                vertex_channels[key[2:]] = obj[key]
+            elif key.startswith('f_'):
+                face_channels[key[2:]] = obj[key]
+        return cls(verts=verts, faces=faces, normals=normals,
+            vertex_channels=vertex_channels, face_channels=face_channels)
+",point_e\util\mesh.py
+save,Save the mesh to a .npz file.,"def save(self, f: Union[str, BinaryIO]):
+    """"""""""""
+    if isinstance(f, str):
+        with open(f, 'wb') as writer:
+            self.save(writer)
+    else:
+        obj_dict = dict(verts=self.verts, faces=self.faces)
+        if self.normals is not None:
+            obj_dict['normals'] = self.normals
+        for k, v in self.vertex_channels.items():
+            obj_dict[f'v_{k}'] = v
+        for k, v in self.face_channels.items():
+            obj_dict[f'f_{k}'] = v
+        np.savez(f, **obj_dict)
+",point_e\util\mesh.py
+has_vertex_colors,,"def has_vertex_colors(self) ->bool:
+    return self.vertex_channels is not None and all(x in self.
+        vertex_channels for x in 'RGB')
+",point_e\util\mesh.py
+write_ply,,"def write_ply(self, raw_f: BinaryIO):
+    write_ply(raw_f, coords=self.verts, rgb=np.stack([self.vertex_channels[
+        x] for x in 'RGB'], axis=1) if self.has_vertex_colors() else None,
+        faces=self.faces)
+",point_e\util\mesh.py
+marching_cubes_mesh,"Run marching cubes on the SDF predicted from a point cloud to produce a
+mesh representing the 3D surface.
+
+:param pc: the point cloud to apply marching cubes to.
+:param model: the model to use to predict SDF values.
+:param grid_size: the number of samples along each axis. A total of
+                  grid_size**3 function evaluations are performed.
+:param side_length: the size of the cube containing the model, which is
+                    assumed to be centered at the origin.
+:param fill_vertex_channels: if True, use the nearest neighbor of each mesh
+                             vertex in the point cloud to compute vertex
+                             data (e.g. colors).","def marching_cubes_mesh(pc: PointCloud, model: PointCloudSDFModel,
+    batch_size: int=4096, grid_size: int=128, side_length: float=1.02,
+    fill_vertex_channels: bool=True, progress: bool=False) ->TriMesh:
+    """"""""""""
+    voxel_size = side_length / (grid_size - 1)
+    min_coord = -side_length / 2
+
+    def int_coord_to_float(int_coords: torch.Tensor) ->torch.Tensor:
+        return int_coords.float() * voxel_size + min_coord
+    with torch.no_grad():
+        cond = model.encode_point_clouds(torch.from_numpy(pc.coords).
+            permute(1, 0).to(model.device)[None])
+    indices = range(0, grid_size ** 3, batch_size)
+    if progress:
+        indices = tqdm(indices)
+    volume = []
+    for i in indices:
+        indices = torch.arange(i, min(i + batch_size, grid_size ** 3), step
+            =1, dtype=torch.int64, device=model.device)
+        zs = int_coord_to_float(indices % grid_size)
+        ys = int_coord_to_float(torch.div(indices, grid_size, rounding_mode
+            ='trunc') % grid_size)
+        xs = int_coord_to_float(torch.div(indices, grid_size ** 2,
+            rounding_mode='trunc'))
+        coords = torch.stack([xs, ys, zs], dim=0)
+        with torch.no_grad():
+            volume.append(model(coords[None], encoded=cond)[0])
+    volume_np = torch.cat(volume).view(grid_size, grid_size, grid_size).cpu(
+        ).numpy()
+    if np.all(volume_np < 0) or np.all(volume_np > 0):
+        volume_np -= np.mean(volume_np)
+    verts, faces, normals, _ = skimage.measure.marching_cubes(volume=
+        volume_np, level=0, allow_degenerate=False, spacing=(voxel_size,) * 3)
+    old_f1 = faces[:, 0].copy()
+    faces[:, 0] = faces[:, 1]
+    faces[:, 1] = old_f1
+    verts += min_coord
+    return TriMesh(verts=verts, faces=faces, normals=normals,
+        vertex_channels=None if not fill_vertex_channels else
+        _nearest_vertex_channels(pc, verts))
+",point_e\util\pc_to_mesh.py
+_nearest_vertex_channels,,"def _nearest_vertex_channels(pc: PointCloud, verts: np.ndarray) ->Dict[str,
+    np.ndarray]:
+    nearest = pc.nearest_points(verts)
+    return {ch: arr[nearest] for ch, arr in pc.channels.items()}
+",point_e\util\pc_to_mesh.py
+plot_point_cloud,"Render a point cloud as a plot to the given image path.
+
+:param pc: the PointCloud to plot.
+:param image_path: the path to save the image, with a file extension.
+:param color: if True, show the RGB colors from the point cloud.
+:param grid_size: the number of random rotations to render.","def plot_point_cloud(pc: PointCloud, color: bool=True, grid_size: int=1,
+    fixed_bounds: Optional[Tuple[Tuple[float, float, float], Tuple[float,
+    float, float]]]=((-0.75, -0.75, -0.75), (0.75, 0.75, 0.75))):
+    """"""""""""
+    fig = plt.figure(figsize=(8, 8))
+    for i in range(grid_size):
+        for j in range(grid_size):
+            ax = fig.add_subplot(grid_size, grid_size, 1 + j + i *
+                grid_size, projection='3d')
+            color_args = {}
+            if color:
+                color_args['c'] = np.stack([pc.channels['R'], pc.channels[
+                    'G'], pc.channels['B']], axis=-1)
+            c = pc.coords
+            if grid_size > 1:
+                theta = np.pi * 2 * (i * grid_size + j) / grid_size ** 2
+                rotation = np.array([[np.cos(theta), -np.sin(theta), 0.0],
+                    [np.sin(theta), np.cos(theta), 0.0], [0.0, 0.0, 1.0]])
+                c = c @ rotation
+            ax.scatter(c[:, 0], c[:, 1], c[:, 2], **color_args)
+            if fixed_bounds is None:
+                min_point = c.min(0)
+                max_point = c.max(0)
+                size = (max_point - min_point).max() / 2
+                center = (min_point + max_point) / 2
+                ax.set_xlim3d(center[0] - size, center[0] + size)
+                ax.set_ylim3d(center[1] - size, center[1] + size)
+                ax.set_zlim3d(center[2] - size, center[2] + size)
+            else:
+                ax.set_xlim3d(fixed_bounds[0][0], fixed_bounds[1][0])
+                ax.set_ylim3d(fixed_bounds[0][1], fixed_bounds[1][1])
+                ax.set_zlim3d(fixed_bounds[0][2], fixed_bounds[1][2])
+    return fig
+",point_e\util\plotting.py
+write_ply,"Write a PLY file for a mesh or a point cloud.
+
+:param coords: an [N x 3] array of floating point coordinates.
+:param rgb: an [N x 3] array of vertex colors, in the range [0.0, 1.0].
+:param faces: an [N x 3] array of triangles encoded as integer indices.","def write_ply(raw_f: BinaryIO, coords: np.ndarray, rgb: Optional[np.ndarray
+    ]=None, faces: Optional[np.ndarray]=None):
+    """"""""""""
+    with buffered_writer(raw_f) as f:
+        f.write(b'ply\n')
+        f.write(b'format binary_little_endian 1.0\n')
+        f.write(bytes(f'element vertex {len(coords)}\n', 'ascii'))
+        f.write(b'property float x\n')
+        f.write(b'property float y\n')
+        f.write(b'property float z\n')
+        if rgb is not None:
+            f.write(b'property uchar red\n')
+            f.write(b'property uchar green\n')
+            f.write(b'property uchar blue\n')
+        if faces is not None:
+            f.write(bytes(f'element face {len(faces)}\n', 'ascii'))
+            f.write(b'property list uchar int vertex_index\n')
+        f.write(b'end_header\n')
+        if rgb is not None:
+            rgb = (rgb * 255.499).round().astype(int)
+            vertices = [(*coord, *rgb) for coord, rgb in zip(coords.tolist(
+                ), rgb.tolist())]
+            format = struct.Struct('<3f3B')
+            for item in vertices:
+                f.write(format.pack(*item))
+        else:
+            format = struct.Struct('<3f')
+            for vertex in coords.tolist():
+                f.write(format.pack(*vertex))
+        if faces is not None:
+            format = struct.Struct('<B3I')
+            for tri in faces.tolist():
+                f.write(format.pack(len(tri), *tri))
+",point_e\util\ply_util.py
+buffered_writer,,"@contextmanager
+def buffered_writer(raw_f: BinaryIO) ->Iterator[io.BufferedIOBase]:
+    if isinstance(raw_f, io.BufferedIOBase):
+        yield raw_f
+    else:
+        f = io.BufferedWriter(raw_f)
+        yield f
+        f.flush()
+",point_e\util\ply_util.py
+preprocess,,"def preprocess(data, channel):
+    if channel in COLORS:
+        return np.round(data * 255.0)
+    return data
+",point_e\util\point_cloud.py
+load,Load the point cloud from a .npz file.,"@classmethod
+def load(cls, f: Union[str, BinaryIO]) ->'PointCloud':
+    """"""""""""
+    if isinstance(f, str):
+        with open(f, 'rb') as reader:
+            return cls.load(reader)
+    else:
+        obj = np.load(f)
+        keys = list(obj.keys())
+        return PointCloud(coords=obj['coords'], channels={k: obj[k] for k in
+            keys if k != 'coords'})
+",point_e\util\point_cloud.py
+save,Save the point cloud to a .npz file.,"def save(self, f: Union[str, BinaryIO]):
+    """"""""""""
+    if isinstance(f, str):
+        with open(f, 'wb') as writer:
+            self.save(writer)
+    else:
+        np.savez(f, coords=self.coords, **self.channels)
+",point_e\util\point_cloud.py
+write_ply,,"def write_ply(self, raw_f: BinaryIO):
+    write_ply(raw_f, coords=self.coords, rgb=np.stack([self.channels[x] for
+        x in 'RGB'], axis=1) if all(x in self.channels for x in 'RGB') else
+        None)
+",point_e\util\point_cloud.py
+random_sample,"Sample a random subset of this PointCloud.
+
+:param num_points: maximum number of points to sample.
+:param subsample_kwargs: arguments to self.subsample().
+:return: a reduced PointCloud, or self if num_points is not less than
+         the current number of points.","def random_sample(self, num_points: int, **subsample_kwargs) ->'PointCloud':
+    """"""""""""
+    if len(self.coords) <= num_points:
+        return self
+    indices = np.random.choice(len(self.coords), size=(num_points,),
+        replace=False)
+    return self.subsample(indices, **subsample_kwargs)
+",point_e\util\point_cloud.py
+farthest_point_sample,"Sample a subset of the point cloud that is evenly distributed in space.
+
+First, a random point is selected. Then each successive point is chosen
+such that it is furthest from the currently selected points.
+
+The time complexity of this operation is O(NM), where N is the original
+number of points and M is the reduced number. Therefore, performance
+can be improved by randomly subsampling points with random_sample()
+before running farthest_point_sample().
+
+:param num_points: maximum number of points to sample.
+:param init_idx: if specified, the first point to sample.
+:param subsample_kwargs: arguments to self.subsample().
+:return: a reduced PointCloud, or self if num_points is not less than
+         the current number of points.","def farthest_point_sample(self, num_points: int, init_idx: Optional[int]=
+    None, **subsample_kwargs) ->'PointCloud':
+    """"""""""""
+    if len(self.coords) <= num_points:
+        return self
+    init_idx = random.randrange(len(self.coords)
+        ) if init_idx is None else init_idx
+    indices = np.zeros([num_points], dtype=np.int64)
+    indices[0] = init_idx
+    sq_norms = np.sum(self.coords ** 2, axis=-1)
+
+    def compute_dists(idx: int):
+        return sq_norms + sq_norms[idx] - 2 * (self.coords @ self.coords[idx])
+    cur_dists = compute_dists(init_idx)
+    for i in range(1, num_points):
+        idx = np.argmax(cur_dists)
+        indices[i] = idx
+        cur_dists = np.minimum(cur_dists, compute_dists(idx))
+    return self.subsample(indices, **subsample_kwargs)
+",point_e\util\point_cloud.py
+subsample,,"def subsample(self, indices: np.ndarray, average_neighbors: bool=False
+    ) ->'PointCloud':
+    if not average_neighbors:
+        return PointCloud(coords=self.coords[indices], channels={k: v[
+            indices] for k, v in self.channels.items()})
+    new_coords = self.coords[indices]
+    neighbor_indices = PointCloud(coords=new_coords, channels={}
+        ).nearest_points(self.coords)
+    neighbor_indices[indices] = np.arange(len(indices))
+    new_channels = {}
+    for k, v in self.channels.items():
+        v_sum = np.zeros_like(v[:len(indices)])
+        v_count = np.zeros_like(v[:len(indices)])
+        np.add.at(v_sum, neighbor_indices, v)
+        np.add.at(v_count, neighbor_indices, 1)
+        new_channels[k] = v_sum / v_count
+    return PointCloud(coords=new_coords, channels=new_channels)
+",point_e\util\point_cloud.py
+select_channels,,"def select_channels(self, channel_names: List[str]) ->np.ndarray:
+    data = np.stack([preprocess(self.channels[name], name) for name in
+        channel_names], axis=-1)
+    return data
+",point_e\util\point_cloud.py
+nearest_points,"For each point in another set of points, compute the point in this
+pointcloud which is closest.
+
+:param points: an [N x 3] array of points.
+:param batch_size: the number of neighbor distances to compute at once.
+                   Smaller values save memory, while larger values may
+                   make the computation faster.
+:return: an [N] array of indices into self.coords.","def nearest_points(self, points: np.ndarray, batch_size: int=16384
+    ) ->np.ndarray:
+    """"""""""""
+    norms = np.sum(self.coords ** 2, axis=-1)
+    all_indices = []
+    for i in range(0, len(points), batch_size):
+        batch = points[i:i + batch_size]
+        dists = norms + np.sum(batch ** 2, axis=-1)[:, None] - 2 * (batch @
+            self.coords.T)
+        all_indices.append(np.argmin(dists, axis=-1))
+    return np.concatenate(all_indices, axis=0)
+",point_e\util\point_cloud.py
+combine,,"def combine(self, other: 'PointCloud') ->'PointCloud':
+    assert self.channels.keys() == other.channels.keys()
+    return PointCloud(coords=np.concatenate([self.coords, other.coords],
+        axis=0), channels={k: np.concatenate([v, other.channels[k]], axis=0
+        ) for k, v in self.channels.items()})
+",point_e\util\point_cloud.py