//  [python]
def optimize(
        self, train_x, train_y, savedir, warmup_steps=500, num_samples=500,
        target_accept_prob=0.8, max_tree_depth=10, lr=0.1,
        adam_lr=0.01, num_map_iter=1000,
    ):
        self.eval()

        # Pre warmup
        if hasattr(self, "hidden_fns_init"):
            k1 = rbf_kernel(train_x, train_x, self.ls1, self.os1)
            hidden_fns_init = torch.linalg.solve_triangular(
                self.hidden_fns_init, psd_safe_cholesky(k1), upper=False
            ).transpose(-1, -2)
            k2 = add_rbf_kernel(self.hidden_fns_init, self.hidden_fns_init, self.ls2, self.os2)
            final_fns_init = torch.linalg.solve_triangular(
                self.final_fns_init.unsqueeze(-1), psd_safe_cholesky(k2), upper=False
            ).transpose(-1, -2)

            initial_params = {
                "hidden_fns": hidden_fns_init,
                "final_fns": final_fns_init,
            }

        else:
            hidden_fns_param = torch.nn.Parameter(torch.randn(self.width, train_x.size(-2), dtype=train_x.dtype, device=train_x.device))
            final_fns_param = torch.nn.Parameter(torch.randn(1, train_x.size(-2), dtype=train_x.dtype, device=train_x.device))

            def guide(inputs, targets=None, train_x=None):
                with pyro.plate("hidden_fns_plate", self.width, dim=-1):
                    hidden_dist = pyro.distributions.Delta(hidden_fns_param).to_event(1)
                    hidden_fns = pyro.sample("hidden_fns", hidden_dist).transpose(-1, -2)
                with pyro.plate("output_fns_plate", 1, dim=-1):
                    final_dist = pyro.distributions.Delta(final_fns_param).to_event(1)
                    pyro.sample("final_fns", final_dist).transpose(-1, -2)

            optimizer = torch.optim.Adam(
                lr=adam_lr, params=[hidden_fns_param, final_fns_param],
            )
            elbo = pyro.infer.Trace_ELBO()
            iterator = tqdm.tqdm(range(num_map_iter), desc="Training", disable=os.getenv("PROGBAR"))
            for _ in iterator:
                optimizer.zero_grad()
                loss = elbo.differentiable_loss(self, guide, train_x, train_y)
                loss.backward()
                optimizer.step()
                iterator.set_postfix(loss=loss.item())

            initial_params = {
                "hidden_fns": hidden_fns_param.detach(),
                "final_fns": final_fns_param.detach(),
            }

        # MCMC chain
        if num_samples > 0:
            kernel = pyro.infer.NUTS(
                self,
                step_size=(lr / len(train_x)),
                adapt_mass_matrix=True,
                adapt_step_size=True,
                max_plate_nesting=1,
                target_accept_prob=target_accept_prob,
                max_tree_depth=max_tree_depth,
            )
            mcmc = pyro.infer.MCMC(
                kernel,
                num_samples=num_samples,
                warmup_steps=warmup_steps,
                initial_params=initial_params,
                num_chains=1,
                disable_progbar=(os.getenv("PROGBAR")),
            )
            mcmc.run(train_x, train_y)
            torch.save(mcmc.get_samples(), os.path.join(savedir, "mcmc.pth"))
        else:
            torch.save(
                dict((key, val.unsqueeze(0)) for key, val in initial_params.items()),
                os.path.join(savedir, "mcmc.pth"),
            )

        return self