# Path: flaxformer/architectures/longt5/tensor_utils.py
# def pad_to_multiple(array: Array,
#                     factor: int,
#                     axis: int,
#                     mode: Optional[str] = 'constant',
#                     constant_values=0) -> Array:
# def split_into_blocks(array: Array,
#                       block_len: int,
#                       axis: int,
#                       pad_value=0) -> Array:
# def concat_3_blocks(blocked_seq: Array,
#                     block_axis: int,
#                     seq_axis: int,
#                     pad_value=0) -> Array:
# def concat_3_blocks_one_hot(blocked_seq: Array,
#                             block_axis: int,
#                             seq_axis: int,
#                             pad_value=0) -> Array:
# def make_3block_local_att_mask(block_len: int,
#                                input_mask: Array,
#                                segment_ids: Optional[Array] = None,
#                                use_full_block_att: bool = False,
#                                use_causal_mask: bool = False) -> Array:
# def make_3block_relative_position(block_len: int) -> np.ndarray:
# def constant_init(value, dtype=jnp.float32):
#   def init(unused_key, shape, dtype=dtype):


from absl.testing import absltest
from absl.testing import parameterized
from flaxformer.architectures.longt5 import tensor_utils
import jax.numpy as jnp
import numpy as np
# Copyright 2022 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://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.

"""Tests for tensor_utils."""




class TensorUtilsTest(parameterized.TestCase):

  def test_pad_to_multiple_1d(self):
    array = np.arange(3) + 1

    np.testing.assert_array_equal(

# Path: flaxformer/architectures/longt5/tensor_utils.py
# def pad_to_multiple(array: Array,
#                     factor: int,
#                     axis: int,
#                     mode: Optional[str] = 'constant',
#                     constant_values=0) -> Array:
# def split_into_blocks(array: Array,
#                       block_len: int,
#                       axis: int,
#                       pad_value=0) -> Array:
# def concat_3_blocks(blocked_seq: Array,
#                     block_axis: int,
#                     seq_axis: int,
#                     pad_value=0) -> Array:
# def concat_3_blocks_one_hot(blocked_seq: Array,
#                             block_axis: int,
#                             seq_axis: int,
#                             pad_value=0) -> Array:
# def make_3block_local_att_mask(block_len: int,
#                                input_mask: Array,
#                                segment_ids: Optional[Array] = None,
#                                use_full_block_att: bool = False,
#                                use_causal_mask: bool = False) -> Array:
# def make_3block_relative_position(block_len: int) -> np.ndarray:
# def constant_init(value, dtype=jnp.float32):
#   def init(unused_key, shape, dtype=dtype):


# Path: flaxformer/architectures/longt5/tensor_utils_test.py
from absl.testing import absltest
from absl.testing import parameterized
from flaxformer.architectures.longt5 import tensor_utils
import jax.numpy as jnp
import numpy as np

# Copyright 2022 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://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.

"""Tests for tensor_utils."""




class TensorUtilsTest(parameterized.TestCase):

  def test_pad_to_multiple_1d(self):
    array = np.arange(3) + 1

    np.testing.assert_array_equal(

        tensor_utils.pad_to_multiple(array, factor=3, axis=0), [1, 2, 3])