Access module function from outside

[braniii]

Hi there,
I am trying to parametrize an energy-based network, which works as intended. My issue is, that I would like to access the ScoreEBMLP.energy after training the network., something like score_model.energy(x, t). Can anyone help me how this can be achieved? Thx for your help.

Here is a short example of my code.

from flax import linen as nn

class ScoreEBMLP(nn.Module):
    """Energy-Based MLP."""
    features: Sequence[int]

    def energy(self, x, t):
        x = np.concatenate((x, t))

        for idx, feature in enumerate(self.features):
            x = nn.Dense(feature)(x)
            if idx < len(self.features) - 1:
                x = nn.swish(x)

        return x.sum()

    @nn.compact
    def __call__(self, x, t):
        return jax.grad(self.energy, argnums=0)(x, t)