prob_model.load_state("../swag_checkpoints/2023-07-25 14:59:40.237855/checkpoint_18600/checkpoint")
state = prob_model.posterior.state.get()
# SWAGState(step=array(18600, dtype=int32), apply_fn=None, params=FrozenDict({
#    model: {
 #       params: {
  #          dfe_subnet: {
   #             BatchNorm_0: {
   #                 bias: array([-0.13133389, -0.14736553, -0.14047779, -0.12409671, -0.11933165,
     #                      -0.16984864, -0.13965459, -0.07937623, -0.11898279, -0.1386996 ,
        #                   -0.13736989, -0.11246286, -0.15424594, -0.10375523, -0.10800011,
           #                -0.14000903, -0.15316793, -0.13276398, -0.11146024, -0.16203304,
              #             -0.14830959, -0.13227627, -0.11291285, -0.11979104, -0.08990214,
                 #          -0.13557586, -0.15480955, -0.17320064, -0.14736709, -0.12703426, ...
state.mean
# array(-0.01478862, dtype=float32)

    207 if state.mutable is not None and inputs_loader is None and inputs is None:
    208     raise ValueError(
    209         "The posterior state contains mutable objects. Please pass `inputs_loader` or `inputs`."
    210     )
--> 212 n_params = len(state.mean) # TypeError: len() of unsized object
    213 rank = state.dev.shape[-1]
    214 which_params = decode_encoded_tuple_of_lists_of_strings_to_array(
    215     state._encoded_which_params
    216 )

# // Only differ by posterior_approximator method 
map_prob_model
laplace_prob_model

checkpoint = "/path/to/map/checkpoint"

# // Validation accuracy of MAP model at checkpoint is as expected...
map_prob_model.load_state("/path/to/map/checkpoint")
map_out = map_prob_model.predictive.mean(val_loader.to_inputs_loader())
(map_out.argmax(axis=-1) == val_loader.to_array_targets()).sum() / val_loader.size
# '0.67'

from fortuna.metric.classification import accuracy
from fortuna.prob_model import FitCheckpointer, FitConfig, FitMonitor, FitOptimizer

optimizer = FitOptimizer(n_epochs=main_epochs)
monitor = FitMonitor(
    metrics=(accuracy,),
    eval_every_n_epochs=1,
)
checkpointer = FitCheckpointer(
    save_checkpoint_dir=main_save_dir,
   # // Start training from the MAP checkpoint
    restore_checkpoint_path="/path/to/map/checkpoint/",
    keep_top_n_checkpoints=2,
)
config = FitConfig(checkpointer=checkpointer, monitor=monitor)
laplace_status = laplace_prob_model.train(
    fit_config=config,
    train_data_loader=train_loader,
    val_data_loader=val_loader,
)
# // Validation accuracy is NOT as expected...
laplace_out = laplace_prob_model.predictive.mean(val_loader.to_inputs_loader())
(laplace_out.argmax(axis=-1) == val_loader.to_array_targets()).sum() / val_loader.size
# '0.11'

import flax.linen as nn
from fortuna.prob_model import ProbClassifier
from fortuna.data import DataLoader 
from fortuna.prob_model import FitConfig

class CNN(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Conv(features=32, kernel_size=(3, 3))(x)
        x = nn.relu(x)
        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = nn.Conv(features=64, kernel_size=(3, 3))(x)
        x = nn.relu(x)
        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = x.reshape((x.shape[0], -1))  # flatten
        x = nn.Dense(features=256)(x)
        x = nn.relu(x)
        x = nn.Dense(features=10)(x)
        x = nn.log_softmax(x)
        return x
    
prob_model = ProbClassifier(model=CNN())


x = jnp.zeros((5, 64, 64, 10))
y = jnp.ones((5,))

train_loader = DataLoader.from_array_data(
    data=(x,y), batch_size=1
)

prob_model.train(
    train_loader,
    fit_config=FitConfig(
    )
)
# RAISES ERROR

output_dim = 10
prob_model = ProbClassifier(
    model=LeNet5(output_dim=output_dim),
    posterior_approximator=LaplacePosteriorApproximator(),
)

status = prob_model.train(
    train_data_loader=train_data_loader,
    val_data_loader=val_data_loader,
    calib_data_loader=val_data_loader,
    fit_config=FitConfig(
        optimizer=FitOptimizer(freeze_fun=lambda path, val: "trainable" if "output_subnet" in path else "frozen")
    ),
    map_fit_config=FitConfig(
        monitor=FitMonitor(early_stopping_patience=2, metrics=(accuracy,)),
        optimizer=FitOptimizer()
    ),
    calib_config=CalibConfig(monitor=CalibMonitor(early_stopping_patience=2))
)

from typing import Any
import flax.linen as nn
import jax.numpy as jnp
import jax

act = jax.nn.swish


class AlexNet(nn.Module):
    output_dim: int
    dtype: Any = jnp.float32
    """
    An AlexNet model for Cifar10.
    """

    def setup(self):
        self.hidden_layers = AlexNetHiddenLayers(dtype=self.dtype)
        self.last_layer = AlexNetLastLayer(output_dim=self.output_dim, dtype=self.dtype)

    def __call__(self, x: jnp.ndarray, train: bool = True) -> jnp.ndarray:
        x = self.hidden_layers(x, train)
        x = self.last_layer(x, train)
        return x


class AlexNetHiddenLayers(nn.Module):
    dtype: Any = jnp.float32
    """
    Hidden Convolutional layers of AlexNet model
    """

    @nn.compact
    def __call__(self, x: jnp.ndarray, train: bool = True):
        # [32, 32, 3]
        x = nn.Conv(features=64, kernel_size=(3,))(x)
        # [32, 32, 64]
        x = act(x)
        x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2))
        # [16, 16, 64]

        x = nn.Conv(features=128, kernel_size=(3,))(x)
        # [16, 16, 128]
        x = act(x)
        x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2))
        # [8, 8, 128]

        x = nn.Conv(features=256, kernel_size=(2,))(x)
        # [8, 8, 256]
        x = act(x)

        x = nn.Conv(features=128, kernel_size=(2,))(x)
        # [8, 8, 128]
        x = act(x)

        x = nn.Conv(features=64, kernel_size=(2,))(x)
        # [8, 8, 64]
        x = act(x)

        x = x.reshape((x.shape[0], -1))
        return x


class AlexNetLastLayer(nn.Module):
    output_dim: int
    dtype: Any = jnp.float32

    @nn.compact
    def __call__(self, x: jnp.ndarray, train: bool = True):
        x = nn.Dense(features=256, dtype=self.dtype)(x)
        x = act(x)
        x = nn.Dense(features=256, dtype=self.dtype)(x)
        x = act(x)
        x = nn.Dense(features=self.output_dim, dtype=self.dtype)(x)
        return x

# // Model
prob_model = ProbClassifier(
    model=AlexNet(output_dim=10), 
    posterior_approximator=LaplacePosteriorApproximator(),
    prior=IsotropicGaussianPrior(log_var=jnp.log(PRIOR_VAR))
)
prob_model.load_state("../sgd_checkpoints/checkpoint_11532/")
test_log_probs = prob_model.predictive.log_prob(data_loader=test_loader)
# RAISES ERROR