Project: s41467-025-59812-0 | Function: mix_prior (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Implements a novel mixed prior distribution combining Normal and Laplace distributions with learnable mixing parameters for Bayesian neural network weights. Creates mixed prior by sampling mixing weight pi from Beta distribution and combining Normal and Laplace components.

**KEY TERMS:** {
  "pi": "mixing weight parameter",
  "normal": "Normal distribution component",
  "laplace": "Laplace distribution component",
  "shape": "layer dimensions"
}

**CONTEXT CODE:** import os

import numpy as np
import jax.nn as jnn
import jax.numpy as jnp
import jax.random as random

import numpyro
import numpyro.distributions as dist


from functools import partial
from numpyro.infer import Predictive, MCMC, NUTS, SVI, Trace_ELBO


from utils import disentangle_uncertainty


def mix_prior(shape1, shape2=None, param_name=None):
    # MASKED: Implements a novel mixed prior distribution combining Normal and Laplace distributions with learnable mixing parameters for Bayesian neural network weights implementation
    raise NotImplementedError("This function is masked for testing purposes.")


class BNN_NUTS(object):
    def __init__(self, input_dims, hidden_dims, output_dims):
        super().__init__()
        self.rng_key, self.rng_key_predict = random.split(random.PRNGKey(666))
        self.hidden_dims = hidden_dims
        self.net = partial(
            self.jax_model, hidden_dims=hidden_dims, output_dims=output_dims)

    # a two-layer bayesian neural network with computational flow

    def jax_model(self, X, Y, hidden_dims, output_dims=1):
        N, input_dims = X.shape

        w1 = mix_prior(input_dims, hidden_dims, param_name="w1")
        b1 = mix_prior(hidden_dims, param_name="b1") / 10.
        assert w1.shape == (input_dims, hidden_dims)
        z1 = jnp.matmul(X, w1) + b1
        z1 = jnn.relu(z1)
        assert z1.shape == (N, hidden_dims)

        # sample final layer of weights and neural network output
        w2 = mix_prior(hidden_dims, output_dims, param_name="w2") / 10.
        b2 = mix_prior(output_dims, param_name="b2") / 10.
        assert w2.shape == (hidden_dims, output_dims)
        z2 = jnp.matmul(z1, w2) + b2
        assert z2.shape == (N, output_dims)

        mean = jnn.sigmoid(z2).squeeze(-1)
        prob = numpyro.deterministic("prob", mean)

        # observe data
        with numpyro.plate("data", N):
            numpyro.sample("Y", dist.Bernoulli(mean), obs=Y)

    def train(self, train_x, train_y):
        kernel = NUTS(self.net)
        mcmc = MCMC(
            kernel,
            num_warmup=200,
            num_samples=100,
            num_chains=4,
            progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
        )
        mcmc.run(self.rng_key, train_x, train_y)
        self.samples = mcmc.get_samples()
        # gelman rubin diagnostics
        # from numpyro.diagnostics import gelman_rubin
        # rhat_values = {k: gelman_rubin(v.reshape(4, -1)) for k, v in self.samples.items()}
        # print(rhat_values)

    def evaluate(self, test_x, num_models=100):
        # obtain predictive posterior of the probability of the bernoulli
        self.predictive_probs = Predictive(
            self.net, self.samples, return_sites=["prob"])
        probs_posterior = self.predictive_probs(
            self.rng_key_predict, test_x, None)["prob"]
        self.save_model("./models/BNN_NUTS.npz")
        predictions = np.array(probs_posterior.T)
        mean, predictive_uncertainty, aleatoric_uncertainty, epistemic_uncertainty = disentangle_uncertainty(
            predictions)
        return mean, predictive_uncertainty, aleatoric_uncertainty, epistemic_uncertainty

    def save_model(self, save_path):
        np.savez(save_path, **self.predictive_probs.posterior_samples)


class BNN_SVI(object):
    def __init__(self, input_dims, hidden_dims, output_dims, lr):
        super().__init__()

        self.rng_key, self.rng_key_predict = random.split(random.PRNGKey(666))
        self.hidden_dims = hidden_dims
        self.optimizer = numpyro.optim.Adam(lr)

    def jax_model(self, X, Y, hidden_dims, output_dims=1):
        N, input_dims = X.shape

        w1 = mix_prior(input_dims, hidden_dims, "w1")
        b1 = mix_prior(hidden_dims, param_name="b1") / 10.

        z1 = jnn.relu(jnp.matmul(X, w1) + b1)

        w2 = mix_prior(hidden_dims, output_dims, "w2") / 10.
        b2 = mix_prior(output_dims, param_name="b2") / 10.
        z2 = jnp.matmul(z1, w2) + b2

        mean = jnn.sigmoid(z2).squeeze(-1)
        prob = numpyro.deterministic("prob", mean)
        with numpyro.plate("data", N):
            numpyro.sample("Y", dist.Bernoulli(mean), obs=Y)

    def guide(self, X, Y, hidden_dims, output_dims):
        N, input_dims = X.shape

        def variational_param(shape, param_name, suffix):


Complete the function implementation replacing "raise NotImplementedError".