# Path: agents/tools/count_weights.py
# def count_weights(scope=None, exclude=None, graph=None):
#   """Count learnable parameters.
# 
#   Args:
#     scope: Restrict the count to a variable scope.
#     exclude: Regex to match variable names to exclude.
#     graph: Operate on a graph other than the current default graph.
# 
#   Returns:
#     Number of learnable parameters as integer.
#   """
#   if scope:
#     scope = scope if scope.endswith('/') else scope + '/'
#   graph = graph or tf.get_default_graph()
#   vars_ = graph.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
#   if scope:
#     vars_ = [var for var in vars_ if var.name.startswith(scope)]
#   if exclude:
#     exclude = re.compile(exclude)
#     vars_ = [var for var in vars_ if not exclude.match(var.name)]
#   shapes = [var.get_shape().as_list() for var in vars_]
#   return int(sum(np.prod(shape) for shape in shapes))


import tensorflow as tf
from agents.tools import count_weights

"""Tests for the weight counting utility."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function




class CountWeightsTest(tf.test.TestCase):

  def test_count_trainable(self):
    tf.Variable(tf.zeros((5, 3)), trainable=True)
    tf.Variable(tf.zeros((1, 1)), trainable=True)
    tf.Variable(tf.zeros((5,)), trainable=True)
    self.assertEqual(15 + 1 + 5, count_weights())

  def test_ignore_non_trainable(self):
    tf.Variable(tf.zeros((5, 3)), trainable=False)
    tf.Variable(tf.zeros((1, 1)), trainable=False)
    tf.Variable(tf.zeros((5,)), trainable=False)
    self.assertEqual(0, count_weights())

  def test_trainable_and_non_trainable(self):
    tf.Variable(tf.zeros((5, 3)), trainable=True)
    tf.Variable(tf.zeros((8, 2)), trainable=False)
    tf.Variable(tf.zeros((1, 1)), trainable=True)
    tf.Variable(tf.zeros((5,)), trainable=True)
    tf.Variable(tf.zeros((3, 1)), trainable=False)

# Path: agents/tools/count_weights.py
# def count_weights(scope=None, exclude=None, graph=None):
#   """Count learnable parameters.
# 
#   Args:
#     scope: Restrict the count to a variable scope.
#     exclude: Regex to match variable names to exclude.
#     graph: Operate on a graph other than the current default graph.
# 
#   Returns:
#     Number of learnable parameters as integer.
#   """
#   if scope:
#     scope = scope if scope.endswith('/') else scope + '/'
#   graph = graph or tf.get_default_graph()
#   vars_ = graph.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
#   if scope:
#     vars_ = [var for var in vars_ if var.name.startswith(scope)]
#   if exclude:
#     exclude = re.compile(exclude)
#     vars_ = [var for var in vars_ if not exclude.match(var.name)]
#   shapes = [var.get_shape().as_list() for var in vars_]
#   return int(sum(np.prod(shape) for shape in shapes))


# Path: agents/tools/count_weights_test.py
import tensorflow as tf
from agents.tools import count_weights


"""Tests for the weight counting utility."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function




class CountWeightsTest(tf.test.TestCase):

  def test_count_trainable(self):
    tf.Variable(tf.zeros((5, 3)), trainable=True)
    tf.Variable(tf.zeros((1, 1)), trainable=True)
    tf.Variable(tf.zeros((5,)), trainable=True)
    self.assertEqual(15 + 1 + 5, count_weights())

  def test_ignore_non_trainable(self):
    tf.Variable(tf.zeros((5, 3)), trainable=False)
    tf.Variable(tf.zeros((1, 1)), trainable=False)
    tf.Variable(tf.zeros((5,)), trainable=False)
    self.assertEqual(0, count_weights())

  def test_trainable_and_non_trainable(self):
    tf.Variable(tf.zeros((5, 3)), trainable=True)
    tf.Variable(tf.zeros((8, 2)), trainable=False)
    tf.Variable(tf.zeros((1, 1)), trainable=True)
    tf.Variable(tf.zeros((5,)), trainable=True)
    tf.Variable(tf.zeros((3, 1)), trainable=False)

    self.assertEqual(15 + 1 + 5, count_weights())