# Path: lamb/utils.py
# def ensure_list(x):
# def linear(args, output_size, bias, bias_start=0.0, initializer=None,
#            scope=None):
# def _linear(args, output_size, bias, bias_start=0.0):
# def _find_var_init_param(var_init_params, var_name, key, default):
# def find_initializer(var_init_params, name, default_initializer):
# def linear_v2(args, output_size, bias, rank=None, scope=None,   # pylint: disable=missing-docstring
#               weight_name=_WEIGHTS_VARIABLE_NAME,
#               bias_name=_BIAS_VARIABLE_NAME,
#               var_init_params=None):
# def layer_norm(x, reduction_indices, epsilon=1e-9, gain=None, bias=None,
#                per_element=True, scope=None):
# def get_sparse_variable(name, indices, shape, dtype=None, trainable=True,
#                         initializer=None, partitioner=None, regularizer=None):
# def _random_index(shape):
# def _all_indices(shape):  # pylint: disable=missing-docstring
#   def add_indices(shape, depth, index):
# def sparse_random_indices(ratio, shape):
# def get_optimizer(optimizer_type):
#   def build(learning_rate, config):
# def _extract_dict_from_config(config, prefix, keys):
# def repeat(iterable, n):
# def is_var_in_scope(var, scopes):
# def trainable_vars_in_scope(name):
# def find_var(name, vars_=None):
# def group_vars_by_scope(scopes, vars_=None, log=False):
# def _order_grouped_vars(var_groups):
# def gradients_for_var_group(var_groups, gradients, name):
# def trainable_initial_state(batch_size, state_size, initial_state_init=None):
#   def create_one(i, size):
# def map_nested(fn, x):
# def create_grads(optimizer, loss, scopes, num_expected_missing_gradients=0):
# def clip_gradients_in_scope(grads_and_vars, scope, max_grad_norm):
# def cv_splits(seq, k=10):
# def escape_summary_name(name):
# def summaries_for_trainables():
# def log_scalar_summaries(summary_proto):
# def count_trainables(scopes=('',)):
# def log_trainables(scopes=('',)):
# def round_to_int(x):
# def orthogonal_initializer(gain=1.0, dtype=tf.float32):
#   def initialize(shape, dtype=dtype, partition_info=None, gain=gain):
# def random_mask(shape, k):
# def random_mask2(shape, k):
# def make_low_rank_factorization_initializer(shape, rank):
# def low_rank_factorization(name, shape, rank, initializer=None,
#                            trainable=True, collections=None):
#   def __init__(self, name='serialized', initializer='{}'):
#   def store(self, obj):
#   def retrieve(self):
#   def variables(self):
# def mixture_of_softmaxes(x, k, e, to_logits):
# def expand_tile(tensor, n, name=None):
# def get_batched_variable(name, runtime_batch_size, shape=None,
#                          dtype=tf.float32, initializer=None,
#                          trainable=True):
# def create_mask(lengths, max_length):
# def mask_from_lengths(lengths, max_length=None, dtype=None, name=None):
# def compute_lengths(symbols_list, eos_symbol, name=None,
#                     dtype=tf.int64):
# def seq_softmax_cross_entropy_with_logits(logits, labels, lengths,
#                                           max_length=None, reduce_sum=True,
#                                           name=None):
# def variance_scaling_initializer(scale=2.0, mode='fan_in',
#                                  distribution='truncated_normal',
#                                  mean=0.0, seed=None, dtype=tf.float32):
#   def _initializer(shape, dtype=dtype, partition_info=None):
#   def __init__(self, fn, output_size, reuse=None):
#   def state_size(self):
#   def output_size(self):
#   def zero_state(self, batch_size, dtype):
#   def call(self, inputs, state):
# _BIAS_VARIABLE_NAME = 'bias'
# _WEIGHTS_VARIABLE_NAME = 'kernel'
# class TFSerializer(object):
# class FNCell(tf.nn.rnn_cell.RNNCell):


import re
import six
import tensorflow.compat.v1 as tf
from absl import logging
from lamb import utils
from sonnet.python.modules import base as snt_base
  represents the proportion of entries non-missing entries in the
  tile. The actual connectivity pattern is determined randomly.

  In the future, there may be support for band and block diagonal
  matrices.
  """

  def __init__(self, input_name_and_sizes, output_name_and_sizes,
               var_init_params=None, name='sparse_tiled_linear'):
    super(SparseTiledLinear, self).__init__(
        input_name_and_sizes, output_name_and_sizes,
        var_init_params=var_init_params, name=name)
    self._sparse_indices_cache = {}
    # Cache the SparseTensor instances to avoid the considerable
    # overhead of creating duplicates just to be optimized out.
    self._sparse_variable_cache = {}

  def _find_or_create_sparse_indices(self, name, shape):
    ratio = self._find_var_init_param(name, 'sparsity_ratio', None)
    assert ratio, 'sparsity_ratio must be specified.'
    sharing_key = self._find_var_init_param(name, 'sparse_indices_sharing_key',
                                            ':name:')
    if sharing_key == ':name:':
      key = name
    if sharing_key == ':shape:':
      sharing_key = shape
    key = (sharing_key, ratio)
    if key not in self._sparse_indices_cache:
      logging.info('Creating sparse indices for %s%r with key %r.',
                   name, shape, key)

# Path: lamb/utils.py
# def ensure_list(x):
# def linear(args, output_size, bias, bias_start=0.0, initializer=None,
#            scope=None):
# def _linear(args, output_size, bias, bias_start=0.0):
# def _find_var_init_param(var_init_params, var_name, key, default):
# def find_initializer(var_init_params, name, default_initializer):
# def linear_v2(args, output_size, bias, rank=None, scope=None,   # pylint: disable=missing-docstring
#               weight_name=_WEIGHTS_VARIABLE_NAME,
#               bias_name=_BIAS_VARIABLE_NAME,
#               var_init_params=None):
# def layer_norm(x, reduction_indices, epsilon=1e-9, gain=None, bias=None,
#                per_element=True, scope=None):
# def get_sparse_variable(name, indices, shape, dtype=None, trainable=True,
#                         initializer=None, partitioner=None, regularizer=None):
# def _random_index(shape):
# def _all_indices(shape):  # pylint: disable=missing-docstring
#   def add_indices(shape, depth, index):
# def sparse_random_indices(ratio, shape):
# def get_optimizer(optimizer_type):
#   def build(learning_rate, config):
# def _extract_dict_from_config(config, prefix, keys):
# def repeat(iterable, n):
# def is_var_in_scope(var, scopes):
# def trainable_vars_in_scope(name):
# def find_var(name, vars_=None):
# def group_vars_by_scope(scopes, vars_=None, log=False):
# def _order_grouped_vars(var_groups):
# def gradients_for_var_group(var_groups, gradients, name):
# def trainable_initial_state(batch_size, state_size, initial_state_init=None):
#   def create_one(i, size):
# def map_nested(fn, x):
# def create_grads(optimizer, loss, scopes, num_expected_missing_gradients=0):
# def clip_gradients_in_scope(grads_and_vars, scope, max_grad_norm):
# def cv_splits(seq, k=10):
# def escape_summary_name(name):
# def summaries_for_trainables():
# def log_scalar_summaries(summary_proto):
# def count_trainables(scopes=('',)):
# def log_trainables(scopes=('',)):
# def round_to_int(x):
# def orthogonal_initializer(gain=1.0, dtype=tf.float32):
#   def initialize(shape, dtype=dtype, partition_info=None, gain=gain):
# def random_mask(shape, k):
# def random_mask2(shape, k):
# def make_low_rank_factorization_initializer(shape, rank):
# def low_rank_factorization(name, shape, rank, initializer=None,
#                            trainable=True, collections=None):
#   def __init__(self, name='serialized', initializer='{}'):
#   def store(self, obj):
#   def retrieve(self):
#   def variables(self):
# def mixture_of_softmaxes(x, k, e, to_logits):
# def expand_tile(tensor, n, name=None):
# def get_batched_variable(name, runtime_batch_size, shape=None,
#                          dtype=tf.float32, initializer=None,
#                          trainable=True):
# def create_mask(lengths, max_length):
# def mask_from_lengths(lengths, max_length=None, dtype=None, name=None):
# def compute_lengths(symbols_list, eos_symbol, name=None,
#                     dtype=tf.int64):
# def seq_softmax_cross_entropy_with_logits(logits, labels, lengths,
#                                           max_length=None, reduce_sum=True,
#                                           name=None):
# def variance_scaling_initializer(scale=2.0, mode='fan_in',
#                                  distribution='truncated_normal',
#                                  mean=0.0, seed=None, dtype=tf.float32):
#   def _initializer(shape, dtype=dtype, partition_info=None):
#   def __init__(self, fn, output_size, reuse=None):
#   def state_size(self):
#   def output_size(self):
#   def zero_state(self, batch_size, dtype):
#   def call(self, inputs, state):
# _BIAS_VARIABLE_NAME = 'bias'
# _WEIGHTS_VARIABLE_NAME = 'kernel'
# class TFSerializer(object):
# class FNCell(tf.nn.rnn_cell.RNNCell):


# Path: lamb/tiled_linear.py
import re
import six
import tensorflow.compat.v1 as tf
from absl import logging
from lamb import utils
from sonnet.python.modules import base as snt_base

  represents the proportion of entries non-missing entries in the
  tile. The actual connectivity pattern is determined randomly.

  In the future, there may be support for band and block diagonal
  matrices.
  """

  def __init__(self, input_name_and_sizes, output_name_and_sizes,
               var_init_params=None, name='sparse_tiled_linear'):
    super(SparseTiledLinear, self).__init__(
        input_name_and_sizes, output_name_and_sizes,
        var_init_params=var_init_params, name=name)
    self._sparse_indices_cache = {}
    # Cache the SparseTensor instances to avoid the considerable
    # overhead of creating duplicates just to be optimized out.
    self._sparse_variable_cache = {}

  def _find_or_create_sparse_indices(self, name, shape):
    ratio = self._find_var_init_param(name, 'sparsity_ratio', None)
    assert ratio, 'sparsity_ratio must be specified.'
    sharing_key = self._find_var_init_param(name, 'sparse_indices_sharing_key',
                                            ':name:')
    if sharing_key == ':name:':
      key = name
    if sharing_key == ':shape:':
      sharing_key = shape
    key = (sharing_key, ratio)
    if key not in self._sparse_indices_cache:
      logging.info('Creating sparse indices for %s%r with key %r.',
                   name, shape, key)

      self._sparse_indices_cache[key] = utils.sparse_random_indices(ratio,