# Path: morph_net/framework/matmul_source_op_handler.py
# class MatMulSourceOpHandler(
#     group_lasso_base_op_handler.GroupLassoBaseSourceOpHandler):
#   def _reduce_dims(self, op):
# 
# Path: morph_net/framework/op_regularizer_manager.py
# ITERATION_LIMIT = 1000000
# class OpSlice(collections.namedtuple('OpSlice', ['op', 'slice'])):
# class Slice(collections.namedtuple('Slice', ['start_index', 'size'])):
# class OpRegularizerManager(object):
# class OpGroup(object):
#   def __str__(self):
#   def __str__(self):
#   def __init__(
#       self,
#       output_boundary: List[tf.Operation],
#       op_handler_dict=None,
#       create_grouping_regularizer=grouping_regularizers.MaxGroupingRegularizer,
#       force_group=None,
#       regularizer_blacklist=None,
#       input_boundary: List[tf.Operation] = None,
#       iteration_limit=ITERATION_LIMIT):
#   def ops(self):
#   def get_regularizer(self, op):
#   def create_op_group_for_op_slice(self, op_slice, is_source=True):
#   def group_op_slices(self, op_slices, omit_source_op_slices=None):
#   def slice_op(self, op, sizes):
#   def process_ops(self, ops):
#   def process_ops_last(self, ops):
#   def is_source_op(self, op):
#   def is_passthrough(self, op):
#   def get_op_slices(self, op):
#   def get_op_group(self, op_slice):
#   def _slice_op_slice(self, op_slice, sizes, size_index, size_count,
#                       new_op_slice_group):
#   def _slice_op_with_sizes(self, op, sizes, is_source, is_resliced):
#   def _get_source_slices(self, sizes, op_slices):
#   def _dfs_for_source_ops(self, output_boundary, input_boundary=None):
#   def _force_group_ops(self, force_group):
#   def __init__(self, op_slice=None, op_groups=None, omit_source_op_slices=None):
#   def op_slices(self):
#   def source_op_slices(self):


from absl.testing import parameterized
from morph_net.framework import matmul_source_op_handler
from morph_net.framework import op_regularizer_manager as orm
import tensorflow.compat.v1 as tf
"""Tests for morph_net.framework.matmul_source_op_handler."""

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


class MatmulSourceOpHandlerTest(parameterized.TestCase, tf.test.TestCase):

  @parameterized.named_parameters(('_slice_all', 3), ('_part_slice', 2))
  def testMatMul2D(self, size):
    inputs = tf.zeros((13, 2))
    handler = matmul_source_op_handler.MatMulSourceOpHandler(0.1)

    kernel = tf.constant([[1, 2, 3], [4, 5, 6]], dtype=tf.float32)
    x = tf.matmul(inputs, kernel, transpose_b=False, name='MatMul')
    op_slice = orm.OpSlice(x.op, orm.Slice(0, size))

    transpose_kernel = tf.constant([[1, 4], [2, 5], [3, 6]], dtype=tf.float32)
    x_other = tf.matmul(
        inputs,
        transpose_kernel,
        transpose_b=True,
        name='MatMulTransposedKernel')

# Path: morph_net/framework/matmul_source_op_handler.py
# class MatMulSourceOpHandler(
#     group_lasso_base_op_handler.GroupLassoBaseSourceOpHandler):
#   def _reduce_dims(self, op):
# 
# Path: morph_net/framework/op_regularizer_manager.py
# ITERATION_LIMIT = 1000000
# class OpSlice(collections.namedtuple('OpSlice', ['op', 'slice'])):
# class Slice(collections.namedtuple('Slice', ['start_index', 'size'])):
# class OpRegularizerManager(object):
# class OpGroup(object):
#   def __str__(self):
#   def __str__(self):
#   def __init__(
#       self,
#       output_boundary: List[tf.Operation],
#       op_handler_dict=None,
#       create_grouping_regularizer=grouping_regularizers.MaxGroupingRegularizer,
#       force_group=None,
#       regularizer_blacklist=None,
#       input_boundary: List[tf.Operation] = None,
#       iteration_limit=ITERATION_LIMIT):
#   def ops(self):
#   def get_regularizer(self, op):
#   def create_op_group_for_op_slice(self, op_slice, is_source=True):
#   def group_op_slices(self, op_slices, omit_source_op_slices=None):
#   def slice_op(self, op, sizes):
#   def process_ops(self, ops):
#   def process_ops_last(self, ops):
#   def is_source_op(self, op):
#   def is_passthrough(self, op):
#   def get_op_slices(self, op):
#   def get_op_group(self, op_slice):
#   def _slice_op_slice(self, op_slice, sizes, size_index, size_count,
#                       new_op_slice_group):
#   def _slice_op_with_sizes(self, op, sizes, is_source, is_resliced):
#   def _get_source_slices(self, sizes, op_slices):
#   def _dfs_for_source_ops(self, output_boundary, input_boundary=None):
#   def _force_group_ops(self, force_group):
#   def __init__(self, op_slice=None, op_groups=None, omit_source_op_slices=None):
#   def op_slices(self):
#   def source_op_slices(self):


# Path: morph_net/framework/matmul_source_op_handler_test.py
from absl.testing import parameterized
from morph_net.framework import matmul_source_op_handler
from morph_net.framework import op_regularizer_manager as orm
import tensorflow.compat.v1 as tf

"""Tests for morph_net.framework.matmul_source_op_handler."""

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


class MatmulSourceOpHandlerTest(parameterized.TestCase, tf.test.TestCase):

  @parameterized.named_parameters(('_slice_all', 3), ('_part_slice', 2))
  def testMatMul2D(self, size):
    inputs = tf.zeros((13, 2))
    handler = matmul_source_op_handler.MatMulSourceOpHandler(0.1)

    kernel = tf.constant([[1, 2, 3], [4, 5, 6]], dtype=tf.float32)
    x = tf.matmul(inputs, kernel, transpose_b=False, name='MatMul')
    op_slice = orm.OpSlice(x.op, orm.Slice(0, size))

    transpose_kernel = tf.constant([[1, 4], [2, 5], [3, 6]], dtype=tf.float32)
    x_other = tf.matmul(
        inputs,
        transpose_kernel,
        transpose_b=True,
        name='MatMulTransposedKernel')

    op_slice_other = orm.OpSlice(x_other.op, orm.Slice(0, size))