# Path: ldif/util/tf_util.py
# def assert_shape(tensor, shape, name):
# def log(msg, t):
# def tile_new_axis(t, axis, length):
# def zero_by_mask(mask, vals, replace_with=0.0):
# def remove_element(t, elt, axis):


import numpy as np
import tensorflow as tf
from parameterized import parameterized
from ldif.util import tf_util
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
"""Tests for ldif.util.tf_util."""



# LDIF is an internal package, should be imported last.
# pylint: disable=g-bad-import-order
# pylint: enable=g-bad-import-order

DISTANCE_EPS = 1e-6


class TfUtilTest(tf.test.TestCase):

  @parameterized.expand([('RemoveSecondRow', 1, 0, [[1.0, 2.0]]),
                         ('RemoveFirstRow', 0, 0, [[3.0, 4.0]]),
                         ('RemoveFirstCol', 0, 1, [[2.0], [4.0]]),
                         ('RemoveSecondCol', 1, 1, [[1.0], [3.0]])])
  def testRemoveElement(self, name, elt, axis, expected):
    initial = tf.constant([[1.0, 2.0], [3.0, 4.0]], dtype=tf.float32)

# Path: ldif/util/tf_util.py
# def assert_shape(tensor, shape, name):
# def log(msg, t):
# def tile_new_axis(t, axis, length):
# def zero_by_mask(mask, vals, replace_with=0.0):
# def remove_element(t, elt, axis):


# Path: ldif/util/tf_util_test.py
import numpy as np
import tensorflow as tf
from parameterized import parameterized
from ldif.util import tf_util

# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
"""Tests for ldif.util.tf_util."""



# LDIF is an internal package, should be imported last.
# pylint: disable=g-bad-import-order
# pylint: enable=g-bad-import-order

DISTANCE_EPS = 1e-6


class TfUtilTest(tf.test.TestCase):

  @parameterized.expand([('RemoveSecondRow', 1, 0, [[1.0, 2.0]]),
                         ('RemoveFirstRow', 0, 0, [[3.0, 4.0]]),
                         ('RemoveFirstCol', 0, 1, [[2.0], [4.0]]),
                         ('RemoveSecondCol', 1, 1, [[1.0], [3.0]])])
  def testRemoveElement(self, name, elt, axis, expected):
    initial = tf.constant([[1.0, 2.0], [3.0, 4.0]], dtype=tf.float32)

    removed = tf_util.remove_element(initial, tf.constant([elt],