# 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
        [
            [1, 1, 1, 1, 1],  #
            [1, 1, 1, 0, 0],  #
        ],
        dtype=np.bool_)
    segment_ids = [
        [1, 2, 2, 3, 3],  #
        [1, 1, 2, 0, 0],  #
    ]

    np.testing.assert_array_equal(
        tensor_utils.make_3block_local_att_mask(
            2, input_mask, segment_ids, use_full_block_att=True),
        np.array(
            [
                [
                    [[0, 0, 1, 0, 0, 0], [0, 0, 0, 1, 1, 0]],  #
                    [[0, 1, 1, 0, 0, 0], [0, 0, 0, 1, 1, 0]],  #
                    [[0, 1, 1, 0, 0, 0], [0, 0, 0, 0, 0, 0]],  #
                ],  #
                [
                    [[0, 0, 1, 1, 0, 0], [0, 0, 1, 1, 0, 0]],  #
                    [[0, 0, 1, 0, 0, 0], [0, 0, 0, 0, 0, 0]],  #
                    [[0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]],  #
                ],  #
            ],
            dtype=np.bool_))

  def test_make_3block_relative_position(self):
    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

        [
            [1, 1, 1, 1, 1],  #
            [1, 1, 1, 0, 0],  #
        ],
        dtype=np.bool_)
    segment_ids = [
        [1, 2, 2, 3, 3],  #
        [1, 1, 2, 0, 0],  #
    ]

    np.testing.assert_array_equal(
        tensor_utils.make_3block_local_att_mask(
            2, input_mask, segment_ids, use_full_block_att=True),
        np.array(
            [
                [
                    [[0, 0, 1, 0, 0, 0], [0, 0, 0, 1, 1, 0]],  #
                    [[0, 1, 1, 0, 0, 0], [0, 0, 0, 1, 1, 0]],  #
                    [[0, 1, 1, 0, 0, 0], [0, 0, 0, 0, 0, 0]],  #
                ],  #
                [
                    [[0, 0, 1, 1, 0, 0], [0, 0, 1, 1, 0, 0]],  #
                    [[0, 0, 1, 0, 0, 0], [0, 0, 0, 0, 0, 0]],  #
                    [[0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]],  #
                ],  #
            ],
            dtype=np.bool_))

  def test_make_3block_relative_position(self):
    np.testing.assert_array_equal(

        tensor_utils.make_3block_relative_position(3),