<commit_before>import functools
from copy import copy

import numpy
import jax
import jax.numpy as jnp
from jax import random
from jax.example_libraries import stax  # from jax.example_libraries import stax
from jax.example_libraries import optimizers as optim
from typing import Callable

from jax.example_libraries.optimizers import OptimizerState

from phi import math
from .. import JAX
from ...math._functional import JitFunction


class StaxNet:

    def __init__(self, initialize: Callable, apply: Callable, input_shape: tuple):
        self._initialize = initialize
        self._apply = apply
        self._input_shape = input_shape
        self.parameters = None
        self._tracers = None

    def initialize(self):
        rnd_key = JAX.rnd_key
        JAX.rnd_key, init_key = random.split(rnd_key)
        out_shape, params64 = self._initialize(init_key, input_shape=self._input_shape)
        if math.get_precision() < 64:
            self.parameters = _recursive_to_float32(params64)

    def __call__(self, *args, **kwargs):
        if self._tracers is not None:
            return self._apply(self._tracers, *args)
        else:
            return self._apply(self.parameters, *args)

<commit_msg>Revert "[learning] Import jax.example_libraries"<commit_after>import functools
from copy import copy

import numpy
import jax
import jax.numpy as jnp
from jax import random
from jax.experimental import stax  # from jax.example_libraries import stax
from jax.experimental import optimizers as optim
from typing import Callable

from jax.experimental.optimizers import OptimizerState

from phi import math
from .. import JAX
from ...math._functional import JitFunction


class StaxNet:

    def __init__(self, initialize: Callable, apply: Callable, input_shape: tuple):
        self._initialize = initialize
        self._apply = apply
        self._input_shape = input_shape
        self.parameters = None
        self._tracers = None

    def initialize(self):
        rnd_key = JAX.rnd_key
        JAX.rnd_key, init_key = random.split(rnd_key)
        out_shape, params64 = self._initialize(init_key, input_shape=self._input_shape)
        if math.get_precision() < 64:
            self.parameters = _recursive_to_float32(params64)

    def __call__(self, *args, **kwargs):
        if self._tracers is not None:
            return self._apply(self._tracers, *args)
        else:
            return self._apply(self.parameters, *args)


{'fraction_non_alphanumeric': 0.07300163132137032, 'fraction_numerical': 0.0065252854812398045, 'mean_word_length': 3.1931623931623934, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6470768', 'n_tokens_mistral': 755, 'n_tokens_neox': 726, 'n_words': 209}