# 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)


import math
import flax
import jax.numpy as jnp
from typing import Callable, Mapping, Optional, Sequence, Tuple
from absl import logging
from rigl.experimental.jax.pruning import init
from rigl.experimental.jax.pruning import masked
      inputs: Input data with dimensions (batch, features), if features has more
        than one dimension, it is flattened.
      num_classes: Number of classes in the dataset.
      features: Number of neurons in each layer, and number of layers (given by
        length of sequence) + one layer for softmax.
      train: If model is being evaluated in training mode or not.
      init_fn: Initialization function used for dense layers.
      activation_fn: Activation function to be used for dense layers.
      masks: Masks of the layers in this model, in the same form as module
        params, or None.
      masked_layer_indices: The layer indices of layers in model to be masked.
      dropout_rate: Dropout rate, if 0 then dropout is not used (default).

    Returns:
      A tensor of shape (batch, num_classes), containing the logit output.
    """
    batch_norm = flax.deprecated.nn.BatchNorm.partial(
        use_running_average=not train, momentum=0.99, epsilon=1e-5)

    depth = 1 + len(features)
    masks = masked.generate_model_masks(depth, masks,
                                        masked_layer_indices)

    # If inputs are in image dimensions, flatten image.
    inputs = inputs.reshape(inputs.shape[0], -1)

    for i, feature_num in enumerate(features):
      if f'MaskedModule_{i}' in masks:
        logging.info('Layer %d is masked in model', i)
        mask = masks[f'MaskedModule_{i}']

# 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/models/mnist_fc.py
import math
import flax
import jax.numpy as jnp
from typing import Callable, Mapping, Optional, Sequence, Tuple
from absl import logging
from rigl.experimental.jax.pruning import init
from rigl.experimental.jax.pruning import masked

      inputs: Input data with dimensions (batch, features), if features has more
        than one dimension, it is flattened.
      num_classes: Number of classes in the dataset.
      features: Number of neurons in each layer, and number of layers (given by
        length of sequence) + one layer for softmax.
      train: If model is being evaluated in training mode or not.
      init_fn: Initialization function used for dense layers.
      activation_fn: Activation function to be used for dense layers.
      masks: Masks of the layers in this model, in the same form as module
        params, or None.
      masked_layer_indices: The layer indices of layers in model to be masked.
      dropout_rate: Dropout rate, if 0 then dropout is not used (default).

    Returns:
      A tensor of shape (batch, num_classes), containing the logit output.
    """
    batch_norm = flax.deprecated.nn.BatchNorm.partial(
        use_running_average=not train, momentum=0.99, epsilon=1e-5)

    depth = 1 + len(features)
    masks = masked.generate_model_masks(depth, masks,
                                        masked_layer_indices)

    # If inputs are in image dimensions, flatten image.
    inputs = inputs.reshape(inputs.shape[0], -1)

    for i, feature_num in enumerate(features):
      if f'MaskedModule_{i}' in masks:
        logging.info('Layer %d is masked in model', i)
        mask = masks[f'MaskedModule_{i}']

        inputs = masked.masked(flax.deprecated.nn.Dense, mask)(