# Path: rigl/experimental/jax/pruning/init.py
# def init(rng, shape, dtype=dtype):
#   if mask is None:
#     return base_init(rng, shape, dtype)
# 
#   # Find the ablated neurons in the mask, to determine correct fan_out.
#   neuron_weight_count = jnp.sum(
#       jnp.reshape(mask, (-1, mask.shape[-1])), axis=0)
#   non_zero_neurons = jnp.sum(neuron_weight_count != 0)
# 
#   # Special case of completely ablated weight matrix/layer.
#   if jnp.sum(non_zero_neurons) == 0:
#     print('Empty weight mask!')
#     return jnp.zeros(shape, dtype)
# 
#   # Neurons have different fan_in w/mask, build up initialization per-unit.
#   init_cols = []
#   rng, *split_rngs = jax.random.split(rng, mask.shape[-1] + 1)
#   for i in range(mask.shape[-1]):
#     # Special case of ablated neuron.
#     if neuron_weight_count[i] == 0:
#       init_cols.append(jnp.zeros(shape[:-1] + (1,), dtype))
#       continue
# 
#     # Fake shape of weight matrix with correct fan_in, and fan_out.
#     sparse_shape = (int(neuron_weight_count[i]), int(non_zero_neurons))
# 
#     # Use only the first column of init from initializer, since faked fan_out.
#     init = base_init(split_rngs[i], sparse_shape, dtype)[Ellipsis, 0]
# 
#     # Expand out to full sparse array.
#     expanded_init = jnp.zeros(
#         mask[Ellipsis, i].shape,
#         dtype).flatten().at[jnp.where(mask[Ellipsis, i].flatten() == 1)].set(init)
#     expanded_init = jnp.reshape(expanded_init, mask[Ellipsis, i].shape)
#     init_cols.append(expanded_init[Ellipsis, jnp.newaxis])
# 
#   return jnp.concatenate(init_cols, axis=-1)
# 
# Path: rigl/experimental/jax/pruning/masked.py
# def masked(module, mask):
#   """Convenience function for masking a FLAX module with MaskedModule."""
#   return MaskedModule.partial(wrapped_module=module, mask=mask)


from typing import Any, Mapping, Optional
from absl.testing import absltest
from rigl.experimental.jax.pruning import init
from rigl.experimental.jax.pruning import masked
import flax
import jax
import jax.numpy as jnp
    self._masked_model_sparse_init = flax.deprecated.nn.Model(
        MaskedDenseSparseInit, initial_params)

    mean_init = jnp.mean(
        self._unmasked_model.params['MaskedModule_0']['unmasked']['kernel'])

    stddev_init = jnp.std(
        self._unmasked_model.params['MaskedModule_0']['unmasked']['kernel'])

    mean_sparse_init = jnp.mean(
        self._masked_model_sparse_init.params['MaskedModule_0']['unmasked']
        ['kernel'])

    stddev_sparse_init = jnp.std(
        self._masked_model_sparse_init.params['MaskedModule_0']['unmasked']
        ['kernel'])

    with self.subTest(name='test_cnn_sparse_init_mean'):
      self.assertBetween(mean_sparse_init, mean_init - 2 * stddev_init,
                         mean_init + 2 * stddev_init)

    with self.subTest(name='test_cnn_sparse_init_stddev'):
      self.assertBetween(stddev_sparse_init, 0.5 * stddev_init,
                         1.5 * stddev_init)

  def test_cnn_sparse_init_kaiming(self):
    """Checks kaiming normal sparse initialization for convolutional layer."""
    _, initial_params = MaskedCNN.init_by_shape(self._rng, (self._input_shape,))
    self._unmasked_model = flax.deprecated.nn.Model(MaskedCNN, initial_params)


# Path: rigl/experimental/jax/pruning/init.py
# def init(rng, shape, dtype=dtype):
#   if mask is None:
#     return base_init(rng, shape, dtype)
# 
#   # Find the ablated neurons in the mask, to determine correct fan_out.
#   neuron_weight_count = jnp.sum(
#       jnp.reshape(mask, (-1, mask.shape[-1])), axis=0)
#   non_zero_neurons = jnp.sum(neuron_weight_count != 0)
# 
#   # Special case of completely ablated weight matrix/layer.
#   if jnp.sum(non_zero_neurons) == 0:
#     print('Empty weight mask!')
#     return jnp.zeros(shape, dtype)
# 
#   # Neurons have different fan_in w/mask, build up initialization per-unit.
#   init_cols = []
#   rng, *split_rngs = jax.random.split(rng, mask.shape[-1] + 1)
#   for i in range(mask.shape[-1]):
#     # Special case of ablated neuron.
#     if neuron_weight_count[i] == 0:
#       init_cols.append(jnp.zeros(shape[:-1] + (1,), dtype))
#       continue
# 
#     # Fake shape of weight matrix with correct fan_in, and fan_out.
#     sparse_shape = (int(neuron_weight_count[i]), int(non_zero_neurons))
# 
#     # Use only the first column of init from initializer, since faked fan_out.
#     init = base_init(split_rngs[i], sparse_shape, dtype)[Ellipsis, 0]
# 
#     # Expand out to full sparse array.
#     expanded_init = jnp.zeros(
#         mask[Ellipsis, i].shape,
#         dtype).flatten().at[jnp.where(mask[Ellipsis, i].flatten() == 1)].set(init)
#     expanded_init = jnp.reshape(expanded_init, mask[Ellipsis, i].shape)
#     init_cols.append(expanded_init[Ellipsis, jnp.newaxis])
# 
#   return jnp.concatenate(init_cols, axis=-1)
# 
# Path: rigl/experimental/jax/pruning/masked.py
# def masked(module, mask):
#   """Convenience function for masking a FLAX module with MaskedModule."""
#   return MaskedModule.partial(wrapped_module=module, mask=mask)


# Path: rigl/experimental/jax/pruning/init_test.py
from typing import Any, Mapping, Optional
from absl.testing import absltest
from rigl.experimental.jax.pruning import init
from rigl.experimental.jax.pruning import masked
import flax
import jax
import jax.numpy as jnp

    self._masked_model_sparse_init = flax.deprecated.nn.Model(
        MaskedDenseSparseInit, initial_params)

    mean_init = jnp.mean(
        self._unmasked_model.params['MaskedModule_0']['unmasked']['kernel'])

    stddev_init = jnp.std(
        self._unmasked_model.params['MaskedModule_0']['unmasked']['kernel'])

    mean_sparse_init = jnp.mean(
        self._masked_model_sparse_init.params['MaskedModule_0']['unmasked']
        ['kernel'])

    stddev_sparse_init = jnp.std(
        self._masked_model_sparse_init.params['MaskedModule_0']['unmasked']
        ['kernel'])

    with self.subTest(name='test_cnn_sparse_init_mean'):
      self.assertBetween(mean_sparse_init, mean_init - 2 * stddev_init,
                         mean_init + 2 * stddev_init)

    with self.subTest(name='test_cnn_sparse_init_stddev'):
      self.assertBetween(stddev_sparse_init, 0.5 * stddev_init,
                         1.5 * stddev_init)

  def test_cnn_sparse_init_kaiming(self):
    """Checks kaiming normal sparse initialization for convolutional layer."""
    _, initial_params = MaskedCNN.init_by_shape(self._rng, (self._input_shape,))
    self._unmasked_model = flax.deprecated.nn.Model(MaskedCNN, initial_params)


    mask = masked.simple_mask(self._unmasked_model, jnp.ones,