# pylint: skip-file

"""
Test the external "raw" interface.

"raw" means that types are not converted to pythonic types but are returned as
PureSNMP object instances.
"""

import sys
from datetime import timedelta
from unittest import skipUnless
from unittest.mock import call, patch

import pytest
from x690.types import (
    Integer,
    ObjectIdentifier,
    OctetString,
    Sequence,
    to_bytes,
)

from puresnmp.aio.api.raw import (
    bulkget,
    bulktable,
    bulkwalk,
    get,
    getnext,
    multiget,
    multiset,
    multiwalk,
    set,
    walk,
)
from puresnmp.const import Version
from puresnmp.exc import NoSuchOID, SnmpError
from puresnmp.pdu import BulkGetRequest, GetNextRequest, GetRequest, VarBind
from puresnmp.types import Counter, Gauge, IpAddress, TimeTicks
from puresnmp.util import BulkResult

from . import readbytes, readbytes_multiple
from .asyncmock import AsyncMock

pytestmark = pytest.mark.skipif(
    sys.version_info < (3, 5), reason="requires python3.5"
)

class TestGet:
    @pytest.mark.asyncio
    async def test_get_string(self):
        data =

        expected = OctetString(
            b"Linux d24cf7f36138 4.4.0-28-generic #47-Ubuntu SMP "
            b"Fri Jun 24 10:09:13 UTC 2016 x86_64"
        )
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.return_value = data
            mck().get_request_id.return_value = 0
            result = await get("::1", "private", "1.2.3")
        assert result == expected

    @pytest.mark.asyncio
    async def test_get_oid(self):
        data = readbytes("get_sysoid_01.hex")
        expected = ObjectIdentifier.from_string("1.3.6.1.4.1.8072.3.2.10")
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.return_value = data
            mck().get_request_id.return_value = 0
            result = await get("::1", "private", "1.2.3")
        assert result == expected

    @pytest.mark.asyncio
    async def test_get_multiple_return_binds(self):
        """
        A "GET" response should only return one varbind.
        """
        data = readbytes("get_sysoid_01_error.hex")
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.return_value = data
            mck().get_request_id.return_value = 0
            with pytest.raises(SnmpError, match="varbind"):
                await get("::1", "private", "1.2.3")

    @pytest.mark.asyncio
    async def test_get_non_existing_oid(self):
        """
        A "GET" response on a non-existing OID should raise an appropriate
        exception.
        """
        data = readbytes("get_non_existing.hex")
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.return_value = data
            mck().get_request_id.return_value = 0
            with pytest.raises(NoSuchOID):
                await get("::1", "private", "1.2.3")

class TestWalk:
    @pytest.mark.asyncio
    async def test_walk(self):
        response_1 = readbytes("walk_response_1.hex")
        response_2 = readbytes("walk_response_2.hex")
        response_3 = readbytes("walk_response_3.hex")

        expected = [
            VarBind(
                ObjectIdentifier.from_string("1.3.6.1.2.1.2.2.1.5.1"),
                Gauge(10000000),
            ),
            VarBind(
                ObjectIdentifier.from_string("1.3.6.1.2.1.2.2.1.5.13"),
                Gauge(4294967295),
            ),
        ]

        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.side_effect = [response_1, response_2, response_3]
            mck().get_request_id.return_value = 0
            result = []
            async for x in walk("::1", "public", "1.3.6.1.2.1.2.2.1.5"):
                result.append(x)
        assert result == expected

    @pytest.mark.asyncio
    async def test_walk_multiple_return_binds(self):
        """
        A "WALK" response should only return one varbind.
        """
        data = readbytes("get_sysoid_01_error.hex")
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.return_value = data
            mck().get_request_id.return_value = 0
            with pytest.raises(SnmpError, match="varbind"):
                async for x in walk("::1", "private", "1.2.3"):
                    pass

class TestBulkTable:
    @pytest.mark.asyncio
    async def test_bulktable(self):
        responses = readbytes_multiple("bulktable_response.hex")

        expected = [
            {
                "0": "1",
                "1": Integer(1),
                "2": OctetString(b"lo"),
                "3": Integer(24),
                "4": Integer(65536),
                "5": Gauge(10000000),
                "6": OctetString(b""),
                "7": Integer(1),
                "8": Integer(1),
                "9": TimeTicks(0),
                "10": Counter(172),
                "11": Counter(2),
                "12": Counter(0),
                "13": Counter(0),
                "14": Counter(0),
                "15": Counter(0),
                "16": Counter(172),
                "17": Counter(2),
                "18": Counter(0),
                "19": Counter(0),
                "20": Counter(0),
                "21": Gauge(0),
                "22": ObjectIdentifier(0, 0),
            },
            {
                "0": "4",
                "1": Integer(4),
                "2": OctetString(b"eth0"),
                "3": Integer(6),
                "4": Integer(1500),
                "5": Gauge(4294967295),
                "6": OctetString(b"\x02B\xac\x11\x00\x02"),
                "7": Integer(1),
                "8": Integer(1),
                "9": TimeTicks(0),
                "10": Counter(548760),
                "11": Counter(3888),
                "12": Counter(0),
                "13": Counter(0),
                "14": Counter(0),
                "15": Counter(0),
                "16": Counter(186660),
                "17": Counter(1875),
                "18": Counter(0),
                "19": Counter(0),
                "20": Counter(0),
                "21": Gauge(0),
                "22": ObjectIdentifier(0, 0),
            },
        ]

        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.side_effect = responses
            mck().get_request_id.return_value = 0
            result = []
            table = await bulktable("127.0.0.1", "private", "1.3.6.1.2.1.2.2")
            for row in table:
                result.append(row)
        assert sorted(result, key=lambda x: x["0"]) == expected

class TestMultiGet:
    @pytest.mark.asyncio
    async def test_multiget(self):
        data = readbytes("multiget_response.hex")
        expected = [
            ObjectIdentifier.from_string("1.3.6.1.4.1.8072.3.2.10"),
            OctetString(
                b"Linux 7fbf2f0c363d 4.4.0-28-generic "
                b"#47-Ubuntu SMP Fri Jun 24 10:09:13 "
                b"UTC 2016 x86_64"
            ),
        ]
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.return_value = data
            result = await multiget(
                "::1",
                "private",
                [
                    "1.3.6.1.2.1.1.2.0",
                    "1.3.6.1.2.1.1.1.0",
                ],
            )
        assert result == expected

class TestMultiWalk:

    # TODO (advanced) figure out why this fails on Python 2!
    @skipUnless(
        sys.version_info >= (3, 0),
        "For some reason this fails on Python 2 and I don't know why",
    )
    @pytest.mark.asyncio
    async def test_multi_walk(self):
        response_1 = readbytes("multiwalk_response_1.hex")
        response_2 = readbytes("multiwalk_response_2.hex")
        response_3 = readbytes("multiwalk_response_3.hex")

        expected = [
            VarBind(
                ObjectIdentifier.from_string("1.3.6.1.2.1.2.2.1.1.1"),
                Integer(1),
            ),
            VarBind(
                ObjectIdentifier.from_string("1.3.6.1.2.1.2.2.1.2.1"),
                OctetString(b"lo"),
            ),
            VarBind(
                ObjectIdentifier.from_string("1.3.6.1.2.1.2.2.1.1.78"),
                Integer(78),
            ),
            VarBind(
                ObjectIdentifier.from_string("1.3.6.1.2.1.2.2.1.2.78"),
                OctetString(b"eth0"),
            ),
        ]

        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.side_effect = [response_1, response_2, response_3]
            result = []
            async for x in multiwalk(
                "::1", "public", ["1.3.6.1.2.1.2.2.1.1", "1.3.6.1.2.1.2.2.1.2"]
            ):
                result.append(x)
        # TODO (advanced): should order matter in the following result?
        assert len(result) == len(expected)

    @pytest.mark.asyncio
    async def test_eom(self):
        """
        A test for a walk operation which runs into the endOfMibView marker
        """

        data_generator = readbytes_multiple("x690/multiwalk_endofmibview.hex")

        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.side_effect = data_generator
            result = []
            async for row in multiwalk(
                "::1",
                "public",
                [
                    "1.3.6.1.6.3.16.1.5.2.1.6.6.95.110.111.110.101.95.1",
                ],
            ):
                result.append(row)

        root = "1.3.6.1.6.3.16.1.5.2.1.6.6.95.110.111.110.101.95.1."
        expected = [
            (root + "0", 1),
            (root + "1", 1),
            (root + "2", 1),
        ]

        simplified_result = [
            (str(oid), value.pythonize()) for oid, value in result
        ]

        assert simplified_result == expected

class TestMultiSet:
    @pytest.mark.asyncio
    async def test_multiset(self):
        """
        Test setting multiple OIDs at once.

        NOTE: The OID '1.3.6.1.2.1.1.5.0' below is manually edited for
              unit-testing. It probably has a different type in the real world!
        """
        data = readbytes("multiset_response.hex")
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.return_value = data
            result = await multiset(
                "::1",
                "private",
                [
                    ("1.3.6.1.2.1.1.4.0", OctetString(b"example@nnthink.com")),
                    ("1.3.6.1.2.1.1.5.0", OctetString(b"hello@world.com")),
                ],
            )
        expected = {
            "1.3.6.1.2.1.1.4.0": OctetString(b"hello@world.com"),
            "1.3.6.1.2.1.1.5.0": OctetString(b"hello@world.com"),
        }
        assert result == expected

class TestGetNext:
    @pytest.mark.asyncio
    async def test_get_call_args(self):
        data = readbytes("dummy.hex")  # any dump would do
        packet = Sequence(
            Integer(Version.V2C),
            OctetString("public"),
            GetNextRequest(0, ObjectIdentifier(1, 2, 3)),
        )
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.return_value = data
            await getnext("::1", "public", "1.2.3")
            mck.assert_called_with(timeout=6)
            mck().send.assert_called_with("::1", 161, to_bytes(packet))

    @pytest.mark.asyncio
    async def test_getnext(self):
        data = readbytes("getnext_response.hex")
        expected = VarBind("1.3.6.1.6.3.1.1.6.1.0", Integer(354522558))

        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.return_value = data
            result = await getnext("::1", "private", "1.3.6.1.5")
        assert result == expected

class TestGetBulkGet:
    @pytest.mark.asyncio
    async def test_get_call_args(self):
        data = readbytes("dummy.hex")  # any dump would do
        packet = Sequence(
            Integer(Version.V2C),
            OctetString("public"),
            BulkGetRequest(
                0, 1, 2, ObjectIdentifier(1, 2, 3), ObjectIdentifier(1, 2, 4)
            ),
        )
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.return_value = data
            await bulkget(
                "::1", "public", ["1.2.3"], ["1.2.4"], max_list_size=2
            )
            mck.assert_called_with(timeout=6)
            mck().send.assert_called_with("::1", 161, to_bytes(packet))

    @pytest.mark.asyncio
    async def test_bulkget(self):
        data = readbytes("bulk_get_response.hex")
        expected = BulkResult(
            {
                "1.3.6.1.2.1.1.1.0": OctetString(
                    b"Linux 7e68e60fe303 4.4.0-28-generic "
                    b"#47-Ubuntu SMP Fri Jun 24 10:09:13 UTC 2016 x86_64"
                )
            },
            {
                "1.3.6.1.2.1.3.1.1.1.10.1.172.17.0.1": Integer(10),
                "1.3.6.1.2.1.3.1.1.2.10.1.172.17.0.1": OctetString(
                    b"\x02B\xe2\xc5\x8d\t"
                ),
                "1.3.6.1.2.1.3.1.1.3.10.1.172.17.0.1": IpAddress(
                    b"\xac\x11\x00\x01"
                ),
                "1.3.6.1.2.1.4.1.0": Integer(1),
                "1.3.6.1.2.1.4.3.0": Counter(57),
            },
        )

        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.return_value = data
            result = await bulkget(
                "::1",
                "public",
                ["1.3.6.1.2.1.1.1"],
                ["1.3.6.1.2.1.3.1"],
                max_list_size=5,
            )
        assert result == expected

    @pytest.mark.asyncio
    async def test_eom(self):
        """
        Test a bulg-get operation which runs into the "endOfMibView" marker.
        """

        data = readbytes("x690/bulk_get_eom_response.hex")
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.return_value = data
            result = await bulkget(
                "::1", "public", [], ["1.2.4"], max_list_size=10
            )

        expected_scalars = {}
        assert result.scalars == expected_scalars

        root = "1.3.6.1.6.3.16.1.5.2.1.6.6.95.110.111.110.101.95.1."
        expected_listing = {
            root + "0": Integer(1),
            root + "1": Integer(1),
            root + "2": Integer(1),
        }
        assert result.listing == expected_listing

class TestGetBulkWalk:
    @pytest.mark.asyncio
    async def test_get_call_args(self):
        data = readbytes("dummy.hex")  # any dump would do
        packet = Sequence(
            Integer(Version.V2C),
            OctetString("public"),
            BulkGetRequest(0, 0, 2, ObjectIdentifier(1, 2, 3)),
        )
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.return_value = data

            # we need to wrap this in a list to consume the generator.
            async for x in bulkwalk("::1", "public", ["1.2.3"], bulk_size=2):
                pass
            mck.assert_called_with(timeout=6)
            mck().send.assert_called_with("::1", 161, to_bytes(packet))

    @pytest.mark.asyncio
    async def test_get_call_args_issue_22(self):
        data = readbytes("dummy.hex")  # any dump would do
        packet = Sequence(
            Integer(Version.V2C),
            OctetString("public"),
            BulkGetRequest(0, 0, 2, ObjectIdentifier(1, 2, 3)),
        )
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.return_value = data
            mck().get_request_id.return_value = 0

            with pytest.raises(TypeError, match=r"OIDS.*list"):
                # we need to wrap this in a list to consume the generator.
                async for x in bulkwalk("::1", "public", "1.2.3", bulk_size=2):
                    pass

    @pytest.mark.asyncio
    async def test_bulkwalk(self):
        req1 = readbytes("bulkwalk_request_1.hex")
        req2 = readbytes("bulkwalk_request_2.hex")
        req3 = readbytes("bulkwalk_request_3.hex")

        responses = [
            readbytes("bulkwalk_response_1.hex"),
            readbytes("bulkwalk_response_2.hex"),
            readbytes("bulkwalk_response_3.hex"),
        ]
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.side_effect = responses

            request_ids = [1001613222, 1001613223, 1001613224]
            mck().get_request_id.side_effect = request_ids

            result = []
            async for x in bulkwalk(
                "127.0.0.1", "private", ["1.3.6.1.2.1.2.2"], bulk_size=20
            ):
                result.append(x)

            mck.assert_called_with(timeout=6)
            assert mck().send.mock_calls == [
                call("127.0.0.1", 161, req1),
                call("127.0.0.1", 161, req2),
                call("127.0.0.1", 161, req3),
            ]

            expected = [
                VarBind("1.3.6.1.2.1.2.2.1.1.1", Integer(1)),
                VarBind("1.3.6.1.2.1.2.2.1.1.10", Integer(10)),
                VarBind("1.3.6.1.2.1.2.2.1.2.1", OctetString(b"lo")),
                VarBind("1.3.6.1.2.1.2.2.1.2.10", OctetString(b"eth0")),
                VarBind("1.3.6.1.2.1.2.2.1.3.1", Integer(24)),
                VarBind("1.3.6.1.2.1.2.2.1.3.10", Integer(6)),
                VarBind("1.3.6.1.2.1.2.2.1.4.1", Integer(65536)),
                VarBind("1.3.6.1.2.1.2.2.1.4.10", Integer(1500)),
                VarBind("1.3.6.1.2.1.2.2.1.5.1", Gauge(10000000)),
                VarBind("1.3.6.1.2.1.2.2.1.5.10", Gauge(4294967295)),
                VarBind("1.3.6.1.2.1.2.2.1.6.1", OctetString(b"")),
                VarBind(
                    "1.3.6.1.2.1.2.2.1.6.10",
                    OctetString(b"\x02\x42\xAC\x11\x00\x02"),
                ),
                VarBind("1.3.6.1.2.1.2.2.1.7.1", Integer(1)),
                VarBind("1.3.6.1.2.1.2.2.1.7.10", Integer(1)),
                VarBind("1.3.6.1.2.1.2.2.1.8.1", Integer(1)),
                VarBind("1.3.6.1.2.1.2.2.1.8.10", Integer(1)),
                VarBind("1.3.6.1.2.1.2.2.1.9.1", TimeTicks(0)),
                VarBind("1.3.6.1.2.1.2.2.1.9.10", TimeTicks(0)),
                VarBind("1.3.6.1.2.1.2.2.1.10.1", Counter(172)),
                VarBind("1.3.6.1.2.1.2.2.1.10.10", Counter(60558)),
                VarBind("1.3.6.1.2.1.2.2.1.11.1", Counter(2)),
                VarBind("1.3.6.1.2.1.2.2.1.11.10", Counter(564)),
                VarBind("1.3.6.1.2.1.2.2.1.12.1", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.12.10", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.13.1", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.13.10", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.14.1", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.14.10", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.15.1", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.15.10", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.16.1", Counter(172)),
                VarBind("1.3.6.1.2.1.2.2.1.16.10", Counter(44295)),
                VarBind("1.3.6.1.2.1.2.2.1.17.1", Counter(2)),
                VarBind("1.3.6.1.2.1.2.2.1.17.10", Counter(442)),
                VarBind("1.3.6.1.2.1.2.2.1.18.1", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.18.10", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.19.1", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.19.10", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.20.1", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.20.10", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.21.1", Gauge(0)),
                VarBind("1.3.6.1.2.1.2.2.1.21.10", Gauge(0)),
                VarBind("1.3.6.1.2.1.2.2.1.22.1", ObjectIdentifier(0, 0)),
                VarBind("1.3.6.1.2.1.2.2.1.22.10", ObjectIdentifier(0, 0)),
            ]
        assert result == expected


        data = readbytes("get_sysoid_01_error.hex")
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.return_value = data
            mck().get_request_id.return_value = 0
            with pytest.raises(SnmpError, match="varbind"):
                async for x in walk("::1", "private", "1.2.3"):
                    pass

# pylint: skip-file

"""
Test the external "raw" interface.

"raw" means that types are not converted to pythonic types but are returned as
PureSNMP object instances.
"""

import sys
from datetime import timedelta
from unittest import skipUnless
from unittest.mock import call, patch

import pytest
from x690.types import (
    Integer,
    ObjectIdentifier,
    OctetString,
    Sequence,
    to_bytes,
)

from puresnmp.aio.api.raw import (
    bulkget,
    bulktable,
    bulkwalk,
    get,
    getnext,
    multiget,
    multiset,
    multiwalk,
    set,
    walk,
)
from puresnmp.const import Version
from puresnmp.exc import NoSuchOID, SnmpError
from puresnmp.pdu import BulkGetRequest, GetNextRequest, GetRequest, VarBind
from puresnmp.types import Counter, Gauge, IpAddress, TimeTicks
from puresnmp.util import BulkResult

from . import readbytes, readbytes_multiple
from .asyncmock import AsyncMock

pytestmark = pytest.mark.skipif(
    sys.version_info < (3, 5), reason="requires python3.5"
)

class TestGet:
    @pytest.mark.asyncio
    async def test_get_string(self):
        data = readbytes("get_sysdescr_01.hex")
        expected = OctetString(
            b"Linux d24cf7f36138 4.4.0-28-generic #47-Ubuntu SMP "
            b"Fri Jun 24 10:09:13 UTC 2016 x86_64"
        )
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.return_value = data
            mck().get_request_id.return_value = 0
            result = await get("::1", "private", "1.2.3")
        assert result == expected

    @pytest.mark.asyncio
    async def test_get_oid(self):
        data = readbytes("get_sysoid_01.hex")
        expected = ObjectIdentifier.from_string("1.3.6.1.4.1.8072.3.2.10")
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.return_value = data
            mck().get_request_id.return_value = 0
            result = await get("::1", "private", "1.2.3")
        assert result == expected

    @pytest.mark.asyncio
    async def test_get_multiple_return_binds(self):
        """
        A "GET" response should only return one varbind.
        """
        data = readbytes("get_sysoid_01_error.hex")
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.return_value = data
            mck().get_request_id.return_value = 0
            with pytest.raises(SnmpError, match="varbind"):
                await get("::1", "private", "1.2.3")

    @pytest.mark.asyncio
    async def test_get_non_existing_oid(self):
        """
        A "GET" response on a non-existing OID should raise an appropriate
        exception.
        """
        data = readbytes("get_non_existing.hex")
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.return_value = data
            mck().get_request_id.return_value = 0
            with pytest.raises(NoSuchOID):
                await get("::1", "private", "1.2.3")

class TestWalk:
    @pytest.mark.asyncio
    async def test_walk(self):
        response_1 = readbytes("walk_response_1.hex")
        response_2 = readbytes("walk_response_2.hex")
        response_3 = readbytes("walk_response_3.hex")

        expected = [
            VarBind(
                ObjectIdentifier.from_string("1.3.6.1.2.1.2.2.1.5.1"),
                Gauge(10000000),
            ),
            VarBind(
                ObjectIdentifier.from_string("1.3.6.1.2.1.2.2.1.5.13"),
                Gauge(4294967295),
            ),
        ]

        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.side_effect = [response_1, response_2, response_3]
            mck().get_request_id.return_value = 0
            result = []
            async for x in walk("::1", "public", "1.3.6.1.2.1.2.2.1.5"):
                result.append(x)
        assert result == expected

    @pytest.mark.asyncio
    async def test_walk_multiple_return_binds(self):
        """
        A "WALK" response should only return one varbind.
        """

class TestBulkTable:
    @pytest.mark.asyncio
    async def test_bulktable(self):
        responses = readbytes_multiple("bulktable_response.hex")

        expected = [
            {
                "0": "1",
                "1": Integer(1),
                "2": OctetString(b"lo"),
                "3": Integer(24),
                "4": Integer(65536),
                "5": Gauge(10000000),
                "6": OctetString(b""),
                "7": Integer(1),
                "8": Integer(1),
                "9": TimeTicks(0),
                "10": Counter(172),
                "11": Counter(2),
                "12": Counter(0),
                "13": Counter(0),
                "14": Counter(0),
                "15": Counter(0),
                "16": Counter(172),
                "17": Counter(2),
                "18": Counter(0),
                "19": Counter(0),
                "20": Counter(0),
                "21": Gauge(0),
                "22": ObjectIdentifier(0, 0),
            },
            {
                "0": "4",
                "1": Integer(4),
                "2": OctetString(b"eth0"),
                "3": Integer(6),
                "4": Integer(1500),
                "5": Gauge(4294967295),
                "6": OctetString(b"\x02B\xac\x11\x00\x02"),
                "7": Integer(1),
                "8": Integer(1),
                "9": TimeTicks(0),
                "10": Counter(548760),
                "11": Counter(3888),
                "12": Counter(0),
                "13": Counter(0),
                "14": Counter(0),
                "15": Counter(0),
                "16": Counter(186660),
                "17": Counter(1875),
                "18": Counter(0),
                "19": Counter(0),
                "20": Counter(0),
                "21": Gauge(0),
                "22": ObjectIdentifier(0, 0),
            },
        ]

        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.side_effect = responses
            mck().get_request_id.return_value = 0
            result = []
            table = await bulktable("127.0.0.1", "private", "1.3.6.1.2.1.2.2")
            for row in table:
                result.append(row)
        assert sorted(result, key=lambda x: x["0"]) == expected

class TestMultiGet:
    @pytest.mark.asyncio
    async def test_multiget(self):
        data = readbytes("multiget_response.hex")
        expected = [
            ObjectIdentifier.from_string("1.3.6.1.4.1.8072.3.2.10"),
            OctetString(
                b"Linux 7fbf2f0c363d 4.4.0-28-generic "
                b"#47-Ubuntu SMP Fri Jun 24 10:09:13 "
                b"UTC 2016 x86_64"
            ),
        ]
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.return_value = data
            result = await multiget(
                "::1",
                "private",
                [
                    "1.3.6.1.2.1.1.2.0",
                    "1.3.6.1.2.1.1.1.0",
                ],
            )
        assert result == expected

class TestMultiWalk:

    # TODO (advanced) figure out why this fails on Python 2!
    @skipUnless(
        sys.version_info >= (3, 0),
        "For some reason this fails on Python 2 and I don't know why",
    )
    @pytest.mark.asyncio
    async def test_multi_walk(self):
        response_1 = readbytes("multiwalk_response_1.hex")
        response_2 = readbytes("multiwalk_response_2.hex")
        response_3 = readbytes("multiwalk_response_3.hex")

        expected = [
            VarBind(
                ObjectIdentifier.from_string("1.3.6.1.2.1.2.2.1.1.1"),
                Integer(1),
            ),
            VarBind(
                ObjectIdentifier.from_string("1.3.6.1.2.1.2.2.1.2.1"),
                OctetString(b"lo"),
            ),
            VarBind(
                ObjectIdentifier.from_string("1.3.6.1.2.1.2.2.1.1.78"),
                Integer(78),
            ),
            VarBind(
                ObjectIdentifier.from_string("1.3.6.1.2.1.2.2.1.2.78"),
                OctetString(b"eth0"),
            ),
        ]

        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.side_effect = [response_1, response_2, response_3]
            result = []
            async for x in multiwalk(
                "::1", "public", ["1.3.6.1.2.1.2.2.1.1", "1.3.6.1.2.1.2.2.1.2"]
            ):
                result.append(x)
        # TODO (advanced): should order matter in the following result?
        assert len(result) == len(expected)

    @pytest.mark.asyncio
    async def test_eom(self):
        """
        A test for a walk operation which runs into the endOfMibView marker
        """

        data_generator = readbytes_multiple("x690/multiwalk_endofmibview.hex")

        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.side_effect = data_generator
            result = []
            async for row in multiwalk(
                "::1",
                "public",
                [
                    "1.3.6.1.6.3.16.1.5.2.1.6.6.95.110.111.110.101.95.1",
                ],
            ):
                result.append(row)

        root = "1.3.6.1.6.3.16.1.5.2.1.6.6.95.110.111.110.101.95.1."
        expected = [
            (root + "0", 1),
            (root + "1", 1),
            (root + "2", 1),
        ]

        simplified_result = [
            (str(oid), value.pythonize()) for oid, value in result
        ]

        assert simplified_result == expected

class TestMultiSet:
    @pytest.mark.asyncio
    async def test_multiset(self):
        """
        Test setting multiple OIDs at once.

        NOTE: The OID '1.3.6.1.2.1.1.5.0' below is manually edited for
              unit-testing. It probably has a different type in the real world!
        """
        data = readbytes("multiset_response.hex")
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.return_value = data
            result = await multiset(
                "::1",
                "private",
                [
                    ("1.3.6.1.2.1.1.4.0", OctetString(b"example@nnthink.com")),
                    ("1.3.6.1.2.1.1.5.0", OctetString(b"hello@world.com")),
                ],
            )
        expected = {
            "1.3.6.1.2.1.1.4.0": OctetString(b"hello@world.com"),
            "1.3.6.1.2.1.1.5.0": OctetString(b"hello@world.com"),
        }
        assert result == expected

class TestGetNext:
    @pytest.mark.asyncio
    async def test_get_call_args(self):
        data = readbytes("dummy.hex")  # any dump would do
        packet = Sequence(
            Integer(Version.V2C),
            OctetString("public"),
            GetNextRequest(0, ObjectIdentifier(1, 2, 3)),
        )
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.return_value = data
            await getnext("::1", "public", "1.2.3")
            mck.assert_called_with(timeout=6)
            mck().send.assert_called_with("::1", 161, to_bytes(packet))

    @pytest.mark.asyncio
    async def test_getnext(self):
        data = readbytes("getnext_response.hex")
        expected = VarBind("1.3.6.1.6.3.1.1.6.1.0", Integer(354522558))

        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.return_value = data
            result = await getnext("::1", "private", "1.3.6.1.5")
        assert result == expected

class TestGetBulkGet:
    @pytest.mark.asyncio
    async def test_get_call_args(self):
        data = readbytes("dummy.hex")  # any dump would do
        packet = Sequence(
            Integer(Version.V2C),
            OctetString("public"),
            BulkGetRequest(
                0, 1, 2, ObjectIdentifier(1, 2, 3), ObjectIdentifier(1, 2, 4)
            ),
        )
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.return_value = data
            await bulkget(
                "::1", "public", ["1.2.3"], ["1.2.4"], max_list_size=2
            )
            mck.assert_called_with(timeout=6)
            mck().send.assert_called_with("::1", 161, to_bytes(packet))

    @pytest.mark.asyncio
    async def test_bulkget(self):
        data = readbytes("bulk_get_response.hex")
        expected = BulkResult(
            {
                "1.3.6.1.2.1.1.1.0": OctetString(
                    b"Linux 7e68e60fe303 4.4.0-28-generic "
                    b"#47-Ubuntu SMP Fri Jun 24 10:09:13 UTC 2016 x86_64"
                )
            },
            {
                "1.3.6.1.2.1.3.1.1.1.10.1.172.17.0.1": Integer(10),
                "1.3.6.1.2.1.3.1.1.2.10.1.172.17.0.1": OctetString(
                    b"\x02B\xe2\xc5\x8d\t"
                ),
                "1.3.6.1.2.1.3.1.1.3.10.1.172.17.0.1": IpAddress(
                    b"\xac\x11\x00\x01"
                ),
                "1.3.6.1.2.1.4.1.0": Integer(1),
                "1.3.6.1.2.1.4.3.0": Counter(57),
            },
        )

        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.return_value = data
            result = await bulkget(
                "::1",
                "public",
                ["1.3.6.1.2.1.1.1"],
                ["1.3.6.1.2.1.3.1"],
                max_list_size=5,
            )
        assert result == expected

    @pytest.mark.asyncio
    async def test_eom(self):
        """
        Test a bulg-get operation which runs into the "endOfMibView" marker.
        """

        data = readbytes("x690/bulk_get_eom_response.hex")
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.return_value = data
            result = await bulkget(
                "::1", "public", [], ["1.2.4"], max_list_size=10
            )

        expected_scalars = {}
        assert result.scalars == expected_scalars

        root = "1.3.6.1.6.3.16.1.5.2.1.6.6.95.110.111.110.101.95.1."
        expected_listing = {
            root + "0": Integer(1),
            root + "1": Integer(1),
            root + "2": Integer(1),
        }
        assert result.listing == expected_listing

class TestGetBulkWalk:
    @pytest.mark.asyncio
    async def test_get_call_args(self):
        data = readbytes("dummy.hex")  # any dump would do
        packet = Sequence(
            Integer(Version.V2C),
            OctetString("public"),
            BulkGetRequest(0, 0, 2, ObjectIdentifier(1, 2, 3)),
        )
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().get_request_id.return_value = 0
            mck().send.return_value = data

            # we need to wrap this in a list to consume the generator.
            async for x in bulkwalk("::1", "public", ["1.2.3"], bulk_size=2):
                pass
            mck.assert_called_with(timeout=6)
            mck().send.assert_called_with("::1", 161, to_bytes(packet))

    @pytest.mark.asyncio
    async def test_get_call_args_issue_22(self):
        data = readbytes("dummy.hex")  # any dump would do
        packet = Sequence(
            Integer(Version.V2C),
            OctetString("public"),
            BulkGetRequest(0, 0, 2, ObjectIdentifier(1, 2, 3)),
        )
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.return_value = data
            mck().get_request_id.return_value = 0

            with pytest.raises(TypeError, match=r"OIDS.*list"):
                # we need to wrap this in a list to consume the generator.
                async for x in bulkwalk("::1", "public", "1.2.3", bulk_size=2):
                    pass

    @pytest.mark.asyncio
    async def test_bulkwalk(self):
        req1 = readbytes("bulkwalk_request_1.hex")
        req2 = readbytes("bulkwalk_request_2.hex")
        req3 = readbytes("bulkwalk_request_3.hex")

        responses = [
            readbytes("bulkwalk_response_1.hex"),
            readbytes("bulkwalk_response_2.hex"),
            readbytes("bulkwalk_response_3.hex"),
        ]
        with patch("puresnmp.aio.api.raw.Transport") as mck:
            mck().send = AsyncMock()
            mck().send.side_effect = responses

            request_ids = [1001613222, 1001613223, 1001613224]
            mck().get_request_id.side_effect = request_ids

            result = []
            async for x in bulkwalk(
                "127.0.0.1", "private", ["1.3.6.1.2.1.2.2"], bulk_size=20
            ):
                result.append(x)

            mck.assert_called_with(timeout=6)
            assert mck().send.mock_calls == [
                call("127.0.0.1", 161, req1),
                call("127.0.0.1", 161, req2),
                call("127.0.0.1", 161, req3),
            ]

            expected = [
                VarBind("1.3.6.1.2.1.2.2.1.1.1", Integer(1)),
                VarBind("1.3.6.1.2.1.2.2.1.1.10", Integer(10)),
                VarBind("1.3.6.1.2.1.2.2.1.2.1", OctetString(b"lo")),
                VarBind("1.3.6.1.2.1.2.2.1.2.10", OctetString(b"eth0")),
                VarBind("1.3.6.1.2.1.2.2.1.3.1", Integer(24)),
                VarBind("1.3.6.1.2.1.2.2.1.3.10", Integer(6)),
                VarBind("1.3.6.1.2.1.2.2.1.4.1", Integer(65536)),
                VarBind("1.3.6.1.2.1.2.2.1.4.10", Integer(1500)),
                VarBind("1.3.6.1.2.1.2.2.1.5.1", Gauge(10000000)),
                VarBind("1.3.6.1.2.1.2.2.1.5.10", Gauge(4294967295)),
                VarBind("1.3.6.1.2.1.2.2.1.6.1", OctetString(b"")),
                VarBind(
                    "1.3.6.1.2.1.2.2.1.6.10",
                    OctetString(b"\x02\x42\xAC\x11\x00\x02"),
                ),
                VarBind("1.3.6.1.2.1.2.2.1.7.1", Integer(1)),
                VarBind("1.3.6.1.2.1.2.2.1.7.10", Integer(1)),
                VarBind("1.3.6.1.2.1.2.2.1.8.1", Integer(1)),
                VarBind("1.3.6.1.2.1.2.2.1.8.10", Integer(1)),
                VarBind("1.3.6.1.2.1.2.2.1.9.1", TimeTicks(0)),
                VarBind("1.3.6.1.2.1.2.2.1.9.10", TimeTicks(0)),
                VarBind("1.3.6.1.2.1.2.2.1.10.1", Counter(172)),
                VarBind("1.3.6.1.2.1.2.2.1.10.10", Counter(60558)),
                VarBind("1.3.6.1.2.1.2.2.1.11.1", Counter(2)),
                VarBind("1.3.6.1.2.1.2.2.1.11.10", Counter(564)),
                VarBind("1.3.6.1.2.1.2.2.1.12.1", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.12.10", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.13.1", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.13.10", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.14.1", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.14.10", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.15.1", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.15.10", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.16.1", Counter(172)),
                VarBind("1.3.6.1.2.1.2.2.1.16.10", Counter(44295)),
                VarBind("1.3.6.1.2.1.2.2.1.17.1", Counter(2)),
                VarBind("1.3.6.1.2.1.2.2.1.17.10", Counter(442)),
                VarBind("1.3.6.1.2.1.2.2.1.18.1", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.18.10", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.19.1", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.19.10", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.20.1", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.20.10", Counter(0)),
                VarBind("1.3.6.1.2.1.2.2.1.21.1", Gauge(0)),
                VarBind("1.3.6.1.2.1.2.2.1.21.10", Gauge(0)),
                VarBind("1.3.6.1.2.1.2.2.1.22.1", ObjectIdentifier(0, 0)),
                VarBind("1.3.6.1.2.1.2.2.1.22.10", ObjectIdentifier(0, 0)),
            ]
        assert result == expected


try:
    from setuptools import setup, Extension
except ImportError:
    # Use distutils.core as a fallback.
    # We won't be able to build the Wheel file on Windows.
    from distutils.core import setup, Extension

extensions = []
if os.name == 'nt':
    ext = Extension(
        'asyncio._overlapped', ['overlapped.c'], libraries=['ws2_32'],
    )
    extensions.append(ext)

setup(
    name="asyncio",
    version="0.4.1",

    description="reference implementation of PEP 3156",
    long_description=open("README").read(),
    url="http://example.com",

    classifiers=[
        "Programming Language :: Python",
        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 3.3",
    ],

    packages=["asyncio"],

    ext_modules=extensions,
)


import Speech
from PyObjCTools.TestSupport import TestCase, min_os_level, min_sdk_level
import objc

class TestSFSpeechRecognitionTaskHelper(Speech.NSObject):
    def speechRecognitionTask_didFinishSuccessfully_(self, a, b):
        pass

class TestSFSpeechRecognitionTask(TestCase):

self.assertEqual(Speech.SFSpeechRecognitionTaskStateStarting, 0)
        self.assertEqual(Speech.SFSpeechRecognitionTaskStateRunning, 1)
        self.assertEqual(Speech.SFSpeechRecognitionTaskStateFinishing, 2)
        self.assertEqual(Speech.SFSpeechRecognitionTaskStateCanceling, 3)
        self.assertEqual(Speech.SFSpeechRecognitionTaskStateCompleted, 4)

    @min_sdk_level("10.15")
    def test_protocols(self):
        objc.protocolNamed("SFSpeechRecognitionTaskDelegate")

    def test_methods(self):
        self.assertArgIsBOOL(
            TestSFSpeechRecognitionTaskHelper.speechRecognitionTask_didFinishSuccessfully_,
            1,
        )

    @min_os_level("10.15")
    def test_methods10_15(self):
        self.assertResultIsBOOL(Speech.SFSpeechRecognitionTask.isFinishing)

        self.assertResultIsBOOL(Speech.SFSpeechRecognitionTask.isCancelled)


import Speech
from PyObjCTools.TestSupport import TestCase, min_os_level, min_sdk_level
import objc

class TestSFSpeechRecognitionTaskHelper(Speech.NSObject):
    def speechRecognitionTask_didFinishSuccessfully_(self, a, b):
        pass

class TestSFSpeechRecognitionTask(TestCase):
    def test_constants(self):
        self.assertEqual(Speech.SFSpeechRecognitionTaskStateStarting, 0)
        self.assertEqual(Speech.SFSpeechRecognitionTaskStateRunning, 1)
        self.assertEqual(Speech.SFSpeechRecognitionTaskStateFinishing, 2)
        self.assertEqual(Speech.SFSpeechRecognitionTaskStateCanceling, 3)
        self.assertEqual(Speech.SFSpeechRecognitionTaskStateCompleted, 4)

    @min_sdk_level("10.15")
    def test_protocols(self):

    def test_methods(self):
        self.assertArgIsBOOL(
            TestSFSpeechRecognitionTaskHelper.speechRecognitionTask_didFinishSuccessfully_,
            1,
        )

    @min_os_level("10.15")
    def test_methods10_15(self):
        self.assertResultIsBOOL(Speech.SFSpeechRecognitionTask.isFinishing)

        self.assertResultIsBOOL(Speech.SFSpeechRecognitionTask.isCancelled)


import Speech
from PyObjCTools.TestSupport import TestCase, min_os_level, min_sdk_level
import objc

class TestSFSpeechRecognitionTaskHelper(Speech.NSObject):
    def speechRecognitionTask_didFinishSuccessfully_(self, a, b):
        pass

class TestSFSpeechRecognitionTask(TestCase):
    def test_constants(self):
        self.assertEqual(Speech.SFSpeechRecognitionTaskStateStarting, 0)
        self.assertEqual(Speech.SFSpeechRecognitionTaskStateRunning, 1)
        self.assertEqual(Speech.SFSpeechRecognitionTaskStateFinishing, 2)
        self.assertEqual(Speech.SFSpeechRecognitionTaskStateCanceling, 3)
        self.assertEqual(Speech.SFSpeechRecognitionTaskStateCompleted, 4)

    @min_sdk_level("10.15")
    def test_protocols(self):
        objc.protocolNamed("SFSpeechRecognitionTaskDelegate")

    def test_methods(self):
        self.assertArgIsBOOL(
            TestSFSpeechRecognitionTaskHelper.speechRecognitionTask_didFinishSuccessfully_,
            1,
        )

    @min_os_level("10.15")
    def test_methods10_15(self):
        self


def mask_attn_weights(w):
        # w has shape [batch, heads, dst_sequence, src_sequence], where information flows from src to dst.

import math
import os

import numpy as np
import tensorflow as tf
from pprint import pformat as pps

tf1 = tf.compat.v1

def default_hparams(trainable=True, dtype=tf.float32, scope='model'):
    return {
        'n_vocab': 50257,
        'n_ctx': 1024,
        'n_embd': 768,
        'n_head': 12,
        'n_layer': 12,
        'res_dropout': 0.0,
        'attn_dropout': 0.0,
        'embed_dropout': 0.0,
        'dtype': dtype,
        'trainable': trainable,
        'scope': scope,
        'precision': 'bfloat16' if dtype == tf.bfloat16 else 'float32',
        'scale_by_depth': False,
        'scale_by_in': False,
    }

def shape_list(x):
    """Deal with dynamic shape in tensorflow cleanly."""
    static = x.shape.as_list()
    dynamic = tf.shape(x)
    return [dynamic[i] if s is None else s for i, s in enumerate(static)]

def softmax(x, axis=-1):
    x = x - tf.reduce_max(x, axis=axis, keepdims=True)
    ex = tf.exp(x)
    return ex / tf.reduce_sum(ex, axis=axis, keepdims=True)

def gelu(x):
    return 0.5*x*(1+tf.tanh(np.sqrt(2/np.pi)*(x+0.044715*tf.pow(x, 3))))

def norm(x, scope, *, axis=-1, epsilon=1e-5, params=None):
    """Normalize to mean = 0, std = 1, then do a diagonal affine transform."""
    with tf1.variable_scope(scope):
        n_state = shape_list(x)[-1]
        if params["precision"] == "bfloat16":
            g = tf1.get_variable('g', [n_state], initializer=tf.constant_initializer(1, dtype=tf.bfloat16), dtype=tf.bfloat16)
            b = tf1.get_variable('b', [n_state], initializer=tf.constant_initializer(0, dtype=tf.bfloat16), dtype=tf.bfloat16)
        else:
            g = tf1.get_variable('g', [n_state], initializer=tf.constant_initializer(1))
            b = tf1.get_variable('b', [n_state], initializer=tf.constant_initializer(0))
        u = tf.reduce_mean(x, axis=axis, keepdims=True)
        s = tf.reduce_mean(tf.square(x-u), axis=axis, keepdims=True)
        x = (x - u) * tf.math.rsqrt(s + epsilon)
        x = x*g + b
        return x

def split_states(x, n):
    """Reshape the last dimension of x into [n, x.shape[-1]/n]."""
    *start, u, v = shape_list(x)
    m = u * v
    return tf.reshape(x, start + [n, m//n])

def merge_states(x):
    """Smash the last two dimensions of x into a single dimension."""
    *start, a, b = shape_list(x)
    return tf.reshape(x, start + [a*b])

def conv1d(x, scope, nf, *, w_init_stdev=0.02, params=None, scale=False):
    if params["scale_by_depth"] and scale: # Scale by sqrt(num_layers), only happens at the final projection before a res block output
        w_init_stdev = w_init_stdev * (1. / math.sqrt(params["n_layer"]))
    if params["scale_by_in"]: # Scale by sqrt(num_input_features)
        w_init_stdev = w_init_stdev * (1. / math.sqrt(shape_list(x)[-1]))

    with tf1.variable_scope(scope):
        *start, nx = shape_list(x)
        if params["precision"] == "bfloat16":
            w = tf1.get_variable('w', [1, nx, nf], initializer=tf.random_normal_initializer(stddev=w_init_stdev, dtype=tf.bfloat16), dtype=tf.bfloat16)
            b = tf1.get_variable('b', [nf], initializer=tf.constant_initializer(0, dtype=tf.bfloat16), dtype=tf.bfloat16)
        else:
            w = tf1.get_variable('w', [1, nx, nf], initializer=tf.random_normal_initializer(stddev=w_init_stdev))
            b = tf1.get_variable('b', [nf], initializer=tf.constant_initializer(0))
        c = tf.reshape(tf.matmul(tf.reshape(x, [-1, nx]), tf.reshape(w, [-1, nf]))+b, start+[nf])
        return c

def attention_mask(nd, ns, *, dtype):
    """1's in the lower triangle, counting from the lower right corner.

    Same as tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd), but doesn't produce garbage on TPUs.
    """
    i = tf.range(nd)[:,None]
    j = tf.range(ns)
    m = i >= j - ns + nd
    return tf.cast(m, dtype)

def attn(x, scope, n_state, *, past, params, batch_size, seq_length, train=False):
    assert x.shape.ndims == 2  # Should be [batch*sequence, features]
    assert n_state % params["n_head"] == 0
    *start, hidden_size = shape_list(x)
    num_attention_heads = params["n_head"]
    assert(hidden_size % num_attention_heads == 0)
    size_per_head = hidden_size // num_attention_heads

    if past is not None:
        assert past.shape.ndims == 5  # Should be [batch, 2, heads, sequence, features], where 2 is [k, v]

    def split_heads(x):
        x = tf.reshape(x, [batch_size, seq_length, num_attention_heads, size_per_head])
        x = split_states(x, params["n_head"])
        return tf.transpose(x, [0, 2, 1, 3])

    def merge_heads(x):
        # Reverse of split_heads
        x = tf.transpose(x, [0, 2, 1, 3])
        x = merge_states(x)
        x = tf.reshape(x, [batch_size * seq_length, num_attention_heads * size_per_head])
        return x

_, _, nd, ns = shape_list(w)
        b = attention_mask(nd, ns, dtype=w.dtype)
        b = tf.reshape(b, [1, 1, nd, ns])
        w = w*b - tf.cast(1e10, w.dtype)*(1-b)
        return w

    def multihead_attn(q, k, v):
        # q, k, v have shape [batch, heads, sequence, features]
        w = tf.matmul(q, k, transpose_b=True)
        w = w * tf.math.rsqrt(tf.cast(shape_list(v)[-1], w.dtype))

        w = mask_attn_weights(w)
        w = softmax(w)

        w = dropout(w, params["attn_dropout"], train)

        a = tf.matmul(w, v)
        return a

    with tf1.variable_scope(scope):
        c = conv1d(x, 'c_attn', n_state*3, params=params)
        q, k, v = map(split_heads, tf.split(c, 3, axis=-1))
        present = tf.stack([k, v], axis=1)
        if past is not None:
            pk, pv = tf.unstack(past, axis=1)
            k = tf.concat([pk, k], axis=-2)
            v = tf.concat([pv, v], axis=-2)
        a = multihead_attn(q, k, v)
        a = merge_heads(a)
        a = conv1d(a, 'c_proj', n_state, params=params)
        a = dropout(a, params["res_dropout"], train)
        return a, present

def mlp(x, scope, n_state, *, params, train=False):
    with tf1.variable_scope(scope):
        nx = shape_list(x)[-1]
        h = gelu(conv1d(x, 'c_fc', n_state, params=params))
        h2 = conv1d(h, 'c_proj', nx, params=params, scale=True)
        h2 = dropout(h2, params["res_dropout"], train)
        return h2

def block(x, scope, *, past, params, attn, train=False, **attn_kws):
    with tf1.variable_scope(scope):
        nx = shape_list(x)[-1]
        ln_1 = norm(x, 'ln_1', params=params)
        a, present = attn(ln_1, 'attn', nx, past=past, params=params, train=train, **attn_kws)
        x = x + a
        ln_2 = norm(x, 'ln_2', params=params)
        m = mlp(ln_2, 'mlp', nx*4, params=params, train=train)
        x = x + m
        return x, present

def past_shape(*, params, batch_size=None, sequence=None):
    return [batch_size, params["n_layer"], 2, params["n_head"], sequence, params["n_embd"] // params["n_head"]]

def expand_tile(value, size):
    """Add a new axis of given size."""
    value = tf.convert_to_tensor(value, name='value')
    ndims = value.shape.ndims
    return tf.tile(tf.expand_dims(value, axis=0), [size] + [1]*ndims)

def positions_for(tokens, past_length):
    batch_size = tf.shape(tokens)[0]
    nsteps = tf.shape(tokens)[1]
    return expand_tile(past_length + tf.range(nsteps), batch_size)

def dropout(x, pdrop, train):
    if train and pdrop > 0:
        x = tf.nn.dropout(x, rate=pdrop)
    return x

def _assert_float_dtype(dtype):
    if not dtype.is_floating:
        raise ValueError("Expected floating point type, got %s." % dtype)
    return dtype

def model(X, params, labels=None, past=None, scope='model', reuse=False, train=False):
    with tf1.variable_scope(scope, reuse=reuse):
        results = {}
        batch, sequence = shape_list(X)

        if params["precision"] == "bfloat16":
            wpe = tf1.get_variable('wpe', [params["n_ctx"], params["n_embd"]], # Position encoding
                             initializer=tf.random_normal_initializer(stddev=0.01, dtype=tf.bfloat16), dtype=tf.bfloat16)
            wte = tf1.get_variable('wte', [params["n_vocab"], params["n_embd"]], # Text encoding
                             initializer=tf.random_normal_initializer(stddev=0.02, dtype=tf.bfloat16), dtype=tf.bfloat16)

        else:
            wpe = tf1.get_variable('wpe', [params["n_ctx"], params["n_embd"]], # Position encoding
                                initializer=tf.random_normal_initializer(stddev=0.01))
            wte = tf1.get_variable('wte', [params["n_vocab"], params["n_embd"]], # Text encoding
                                initializer=tf.random_normal_initializer(stddev=0.02))
        past_length = 0 if past is None else tf.shape(past)[-2]

        wpe = dropout(wpe, params["embed_dropout"], train)
        wte = dropout(wte, params["embed_dropout"], train)

        h = tf.gather(wte, X) + tf.gather(wpe, positions_for(X, past_length))

        ## We keep the representation as a 2D tensor to avoid re-shaping it back and
        ## forth from a 3D tensor to a 2D tensor. Re-shapes are normally free on
        ## the GPU/CPU but may not be free on the TPU, so we want to minimize them to
        ## help the optimizer.
        batch_size, seq_length, hidden_size = shape_list(h)
        h = tf.reshape(h, [batch_size * seq_length, hidden_size])

        # Transformer
        presents = []
        activations = []
        pasts = tf.unstack(past, axis=1) if past is not None else [None] * params["n_layer"]
        assert len(pasts) == params["n_layer"]
        checkpoint=False if 'memory_saving_gradients' not in params else params['memory_saving_gradients']
        every = 1 if 'memory_saving_checkpoints' not in params else params['memory_saving_checkpoints']
        for layer, past in enumerate(pasts):
            def block0(x):
              with tf1.variable_scope(scope, reuse=reuse):
                x1, present = block(x, 'h%d' % layer, past=past, params=params, attn=attn, train=train, batch_size=batch, seq_length=sequence)
                presents.append(present)
                return x1
            @tf.custom_gradient
            def block1(input):
              def grad(dy, variables=None):
                # variables contains the tensors used to calculate
                # first, we use stop_gradient to ensure that the
                input0 = tf.stop_gradient(input)
                # then, we use the disconnected input to recalculate
                # the output for this layer.
                output0 = block0(input0)
                # now that we have the output, we need to calculate
                result = tf.gradients(output0, input0, dy)
                if variables != None:
                  paramcount = lambda vs: sum([np.prod(v.shape.as_list()) for v in vs])
                  def logvars(variables, label, print_variables=True):
                    if print_variables:
                      tf.logging.info("%s (%s parameters): %s", label, paramcount(variables), pps(variables))
                    else:
                      tf.logging.info("%s (%s parameters)", label, paramcount(variables))
                    return variables
                  tf.logging.info("---------")
                  logvars(variables, "block1_grad variables for layer %s" % layer)
                  return result, tf.gradients(output0, variables, dy)
                return result
              output = block0(input)
              return output, grad
            if checkpoint and (isinstance(every, int) and layer % every == 0 or layer in every):
                tf.logging.info('checkpointing layer %d', layer)
                tf.add_to_collection('checkpoints', h)
            if bool(int(os.environ.get('GRADIENT_CHECKPOINTING', '0'))):
              h = block1(h)
            else:
              h = block0(h)
            activations.append(h)
        results['present'] = tf.stack(presents, axis=1)
        results['activations'] = activations
        h = norm(h, 'ln_f', params=params)

        h_flat = tf.reshape(h, [batch*sequence, params["n_embd"]])
        logits = tf.matmul(h_flat, wte, transpose_b=True)
        logits = tf.reshape(logits, [batch, sequence, params["n_vocab"]])
        results['logits'] = logits
        return results


with tf1.variable_scope(scope, reuse=reuse):
                x1, present = block(x, 'h%d' % layer, past=past, params=params, attn=attn, train=train, batch_size=batch, seq_length=sequence)
                presents.append(present)
                return x1

import math
import os

import numpy as np
import tensorflow as tf
from pprint import pformat as pps

tf1 = tf.compat.v1

def default_hparams(trainable=True, dtype=tf.float32, scope='model'):
    return {
        'n_vocab': 50257,
        'n_ctx': 1024,
        'n_embd': 768,
        'n_head': 12,
        'n_layer': 12,
        'res_dropout': 0.0,
        'attn_dropout': 0.0,
        'embed_dropout': 0.0,
        'dtype': dtype,
        'trainable': trainable,
        'scope': scope,
        'precision': 'bfloat16' if dtype == tf.bfloat16 else 'float32',
        'scale_by_depth': False,
        'scale_by_in': False,
    }

def shape_list(x):
    """Deal with dynamic shape in tensorflow cleanly."""
    static = x.shape.as_list()
    dynamic = tf.shape(x)
    return [dynamic[i] if s is None else s for i, s in enumerate(static)]

def softmax(x, axis=-1):
    x = x - tf.reduce_max(x, axis=axis, keepdims=True)
    ex = tf.exp(x)
    return ex / tf.reduce_sum(ex, axis=axis, keepdims=True)

def gelu(x):
    return 0.5*x*(1+tf.tanh(np.sqrt(2/np.pi)*(x+0.044715*tf.pow(x, 3))))

def norm(x, scope, *, axis=-1, epsilon=1e-5, params=None):
    """Normalize to mean = 0, std = 1, then do a diagonal affine transform."""
    with tf1.variable_scope(scope):
        n_state = shape_list(x)[-1]
        if params["precision"] == "bfloat16":
            g = tf1.get_variable('g', [n_state], initializer=tf.constant_initializer(1, dtype=tf.bfloat16), dtype=tf.bfloat16)
            b = tf1.get_variable('b', [n_state], initializer=tf.constant_initializer(0, dtype=tf.bfloat16), dtype=tf.bfloat16)
        else:
            g = tf1.get_variable('g', [n_state], initializer=tf.constant_initializer(1))
            b = tf1.get_variable('b', [n_state], initializer=tf.constant_initializer(0))
        u = tf.reduce_mean(x, axis=axis, keepdims=True)
        s = tf.reduce_mean(tf.square(x-u), axis=axis, keepdims=True)
        x = (x - u) * tf.math.rsqrt(s + epsilon)
        x = x*g + b
        return x

def split_states(x, n):
    """Reshape the last dimension of x into [n, x.shape[-1]/n]."""
    *start, u, v = shape_list(x)
    m = u * v
    return tf.reshape(x, start + [n, m//n])

def merge_states(x):
    """Smash the last two dimensions of x into a single dimension."""
    *start, a, b = shape_list(x)
    return tf.reshape(x, start + [a*b])

def conv1d(x, scope, nf, *, w_init_stdev=0.02, params=None, scale=False):
    if params["scale_by_depth"] and scale: # Scale by sqrt(num_layers), only happens at the final projection before a res block output
        w_init_stdev = w_init_stdev * (1. / math.sqrt(params["n_layer"]))
    if params["scale_by_in"]: # Scale by sqrt(num_input_features)
        w_init_stdev = w_init_stdev * (1. / math.sqrt(shape_list(x)[-1]))

    with tf1.variable_scope(scope):
        *start, nx = shape_list(x)
        if params["precision"] == "bfloat16":
            w = tf1.get_variable('w', [1, nx, nf], initializer=tf.random_normal_initializer(stddev=w_init_stdev, dtype=tf.bfloat16), dtype=tf.bfloat16)
            b = tf1.get_variable('b', [nf], initializer=tf.constant_initializer(0, dtype=tf.bfloat16), dtype=tf.bfloat16)
        else:
            w = tf1.get_variable('w', [1, nx, nf], initializer=tf.random_normal_initializer(stddev=w_init_stdev))
            b = tf1.get_variable('b', [nf], initializer=tf.constant_initializer(0))
        c = tf.reshape(tf.matmul(tf.reshape(x, [-1, nx]), tf.reshape(w, [-1, nf]))+b, start+[nf])
        return c

def attention_mask(nd, ns, *, dtype):
    """1's in the lower triangle, counting from the lower right corner.

    Same as tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd), but doesn't produce garbage on TPUs.
    """
    i = tf.range(nd)[:,None]
    j = tf.range(ns)
    m = i >= j - ns + nd
    return tf.cast(m, dtype)

def attn(x, scope, n_state, *, past, params, batch_size, seq_length, train=False):
    assert x.shape.ndims == 2  # Should be [batch*sequence, features]
    assert n_state % params["n_head"] == 0
    *start, hidden_size = shape_list(x)
    num_attention_heads = params["n_head"]
    assert(hidden_size % num_attention_heads == 0)
    size_per_head = hidden_size // num_attention_heads

    if past is not None:
        assert past.shape.ndims == 5  # Should be [batch, 2, heads, sequence, features], where 2 is [k, v]

    def split_heads(x):
        x = tf.reshape(x, [batch_size, seq_length, num_attention_heads, size_per_head])
        x = split_states(x, params["n_head"])
        return tf.transpose(x, [0, 2, 1, 3])

    def merge_heads(x):
        # Reverse of split_heads
        x = tf.transpose(x, [0, 2, 1, 3])
        x = merge_states(x)
        x = tf.reshape(x, [batch_size * seq_length, num_attention_heads * size_per_head])
        return x

    def mask_attn_weights(w):
        # w has shape [batch, heads, dst_sequence, src_sequence], where information flows from src to dst.
        _, _, nd, ns = shape_list(w)
        b = attention_mask(nd, ns, dtype=w.dtype)
        b = tf.reshape(b, [1, 1, nd, ns])
        w = w*b - tf.cast(1e10, w.dtype)*(1-b)
        return w

    def multihead_attn(q, k, v):
        # q, k, v have shape [batch, heads, sequence, features]
        w = tf.matmul(q, k, transpose_b=True)
        w = w * tf.math.rsqrt(tf.cast(shape_list(v)[-1], w.dtype))

        w = mask_attn_weights(w)
        w = softmax(w)

        w = dropout(w, params["attn_dropout"], train)

        a = tf.matmul(w, v)
        return a

    with tf1.variable_scope(scope):
        c = conv1d(x, 'c_attn', n_state*3, params=params)
        q, k, v = map(split_heads, tf.split(c, 3, axis=-1))
        present = tf.stack([k, v], axis=1)
        if past is not None:
            pk, pv = tf.unstack(past, axis=1)
            k = tf.concat([pk, k], axis=-2)
            v = tf.concat([pv, v], axis=-2)
        a = multihead_attn(q, k, v)
        a = merge_heads(a)
        a = conv1d(a, 'c_proj', n_state, params=params)
        a = dropout(a, params["res_dropout"], train)
        return a, present

def mlp(x, scope, n_state, *, params, train=False):
    with tf1.variable_scope(scope):
        nx = shape_list(x)[-1]
        h = gelu(conv1d(x, 'c_fc', n_state, params=params))
        h2 = conv1d(h, 'c_proj', nx, params=params, scale=True)
        h2 = dropout(h2, params["res_dropout"], train)
        return h2

def block(x, scope, *, past, params, attn, train=False, **attn_kws):
    with tf1.variable_scope(scope):
        nx = shape_list(x)[-1]
        ln_1 = norm(x, 'ln_1', params=params)
        a, present = attn(ln_1, 'attn', nx, past=past, params=params, train=train, **attn_kws)
        x = x + a
        ln_2 = norm(x, 'ln_2', params=params)
        m = mlp(ln_2, 'mlp', nx*4, params=params, train=train)
        x = x + m
        return x, present

def past_shape(*, params, batch_size=None, sequence=None):
    return [batch_size, params["n_layer"], 2, params["n_head"], sequence, params["n_embd"] // params["n_head"]]

def expand_tile(value, size):
    """Add a new axis of given size."""
    value = tf.convert_to_tensor(value, name='value')
    ndims = value.shape.ndims
    return tf.tile(tf.expand_dims(value, axis=0), [size] + [1]*ndims)

def positions_for(tokens, past_length):
    batch_size = tf.shape(tokens)[0]
    nsteps = tf.shape(tokens)[1]
    return expand_tile(past_length + tf.range(nsteps), batch_size)

def dropout(x, pdrop, train):
    if train and pdrop > 0:
        x = tf.nn.dropout(x, rate=pdrop)
    return x

def _assert_float_dtype(dtype):
    if not dtype.is_floating:
        raise ValueError("Expected floating point type, got %s." % dtype)
    return dtype

def model(X, params, labels=None, past=None, scope='model', reuse=False, train=False):
    with tf1.variable_scope(scope, reuse=reuse):
        results = {}
        batch, sequence = shape_list(X)

        if params["precision"] == "bfloat16":
            wpe = tf1.get_variable('wpe', [params["n_ctx"], params["n_embd"]], # Position encoding
                             initializer=tf.random_normal_initializer(stddev=0.01, dtype=tf.bfloat16), dtype=tf.bfloat16)
            wte = tf1.get_variable('wte', [params["n_vocab"], params["n_embd"]], # Text encoding
                             initializer=tf.random_normal_initializer(stddev=0.02, dtype=tf.bfloat16), dtype=tf.bfloat16)

        else:
            wpe = tf1.get_variable('wpe', [params["n_ctx"], params["n_embd"]], # Position encoding
                                initializer=tf.random_normal_initializer(stddev=0.01))
            wte = tf1.get_variable('wte', [params["n_vocab"], params["n_embd"]], # Text encoding
                                initializer=tf.random_normal_initializer(stddev=0.02))
        past_length = 0 if past is None else tf.shape(past)[-2]

        wpe = dropout(wpe, params["embed_dropout"], train)
        wte = dropout(wte, params["embed_dropout"], train)

        h = tf.gather(wte, X) + tf.gather(wpe, positions_for(X, past_length))

        ## We keep the representation as a 2D tensor to avoid re-shaping it back and
        ## forth from a 3D tensor to a 2D tensor. Re-shapes are normally free on
        ## the GPU/CPU but may not be free on the TPU, so we want to minimize them to
        ## help the optimizer.
        batch_size, seq_length, hidden_size = shape_list(h)
        h = tf.reshape(h, [batch_size * seq_length, hidden_size])

        # Transformer
        presents = []
        activations = []
        pasts = tf.unstack(past, axis=1) if past is not None else [None] * params["n_layer"]
        assert len(pasts) == params["n_layer"]
        checkpoint=False if 'memory_saving_gradients' not in params else params['memory_saving_gradients']
        every = 1 if 'memory_saving_checkpoints' not in params else params['memory_saving_checkpoints']
        for layer, past in enumerate(pasts):
            def block0(x):

            @tf.custom_gradient
            def block1(input):
              def grad(dy, variables=None):
                # variables contains the tensors used to calculate
                # first, we use stop_gradient to ensure that the
                input0 = tf.stop_gradient(input)
                # then, we use the disconnected input to recalculate
                # the output for this layer.
                output0 = block0(input0)
                # now that we have the output, we need to calculate
                result = tf.gradients(output0, input0, dy)
                if variables != None:
                  paramcount = lambda vs: sum([np.prod(v.shape.as_list()) for v in vs])
                  def logvars(variables, label, print_variables=True):
                    if print_variables:
                      tf.logging.info("%s (%s parameters): %s", label, paramcount(variables), pps(variables))
                    else:
                      tf.logging.info("%s (%s parameters)", label, paramcount(variables))
                    return variables
                  tf.logging.info("---------")
                  logvars(variables, "block1_grad variables for layer %s" % layer)
                  return result, tf.gradients(output0, variables, dy)
                return result
              output = block0(input)
              return output, grad
            if checkpoint and (isinstance(every, int) and layer % every == 0 or layer in every):
                tf.logging.info('checkpointing layer %d', layer)
                tf.add_to_collection('checkpoints', h)
            if bool(int(os.environ.get('GRADIENT_CHECKPOINTING', '0'))):
              h = block1(h)
            else:
              h = block0(h)
            activations.append(h)
        results['present'] = tf.stack(presents, axis=1)
        results['activations'] = activations
        h = norm(h, 'ln_f', params=params)

        h_flat = tf.reshape(h, [batch*sequence, params["n_embd"]])
        logits = tf.matmul(h_flat, wte, transpose_b=True)
        logits = tf.reshape(logits, [batch, sequence, params["n_vocab"]])
        results['logits'] = logits
        return results


import math
import os

import numpy as np
import tensorflow as tf
from pprint import pformat as pps

tf1 = tf.compat.v1

def default_hparams(trainable=True, dtype=tf.float32, scope='model'):
    return {
        'n_vocab': 50257,
        'n_ctx': 1024,
        'n_embd': 768,
        'n_head': 12,
        'n_layer': 12,
        'res_dropout': 0.0,
        'attn_dropout': 0.0,
        'embed_dropout': 0.0,
        'dtype': dtype,
        'trainable': trainable,
        'scope': scope,
        'precision': 'bfloat16' if dtype == tf.bfloat16 else 'float32',
        'scale_by_depth': False,
        'scale_by_in': False,
    }

def shape_list(x):
    """Deal with dynamic shape in tensorflow cleanly."""
    static = x.shape.as_list()
    dynamic = tf.shape(x)
    return [dynamic[i] if s is None else s for i, s in enumerate(static)]

def softmax(x, axis=-1):
    x = x - tf.reduce_max(x, axis=axis, keepdims=True)
    ex = tf.exp(x)
    return ex / tf.reduce_sum(ex, axis=axis, keepdims=True)

def gelu(x):
    return 0.5*x*(1+tf.tanh(np.sqrt(2/np.pi)*(x+0.044715*tf.pow(x, 3))))

def norm(x, scope, *, axis=-1, epsilon=1e-5, params=None):
    """Normalize to mean = 0, std = 1, then do a diagonal affine transform."""
    with tf1.variable_scope(scope):
        n_state = shape_list(x)[-1]
        if params["precision"] == "bfloat16":
            g = tf1.get_variable('g', [n_state], initializer=tf.constant_initializer(1, dtype=tf.bfloat16), dtype=tf.bfloat16)
            b = tf1.get_variable('b', [n_state], initializer=tf.constant_initializer(0, dtype=tf.bfloat16), dtype=tf.bfloat16)
        else:
            g = tf1.get_variable('g', [n_state], initializer=tf.constant_initializer(1))
            b = tf1.get_variable('b', [n_state], initializer=tf.constant_initializer(0))
        u = tf.reduce_mean(x, axis=axis, keepdims=True)
        s = tf.reduce_mean(tf.square(x-u), axis=axis, keepdims=True)
        x = (x - u) * tf.math.rsqrt(s + epsilon)
        x = x*g + b
        return x

def split_states(x, n):
    """Reshape the last dimension of x into [n, x.shape[-1]/n]."""
    *start, u, v = shape_list(x)
    m = u * v
    return tf.reshape(x, start + [n, m//n])

def merge_states(x):
    """Smash the last two dimensions of x into a single dimension."""
    *start, a, b = shape_list(x)
    return tf.reshape(x, start + [a*b])

def conv1d(x, scope, nf, *, w_init_stdev=0.02, params=None, scale=False):
    if params["scale_by_depth"] and scale: # Scale by sqrt(num_layers), only happens at the final projection before a res block output
        w_init_stdev = w_init_stdev * (1. / math.sqrt(params["n_layer"]))
    if params["scale_by_in"]: # Scale by sqrt(num_input_features)
        w_init_stdev = w_init_stdev * (1. / math.sqrt(shape_list(x)[-1]))

    with tf1.variable_scope(scope):
        *start, nx = shape_list(x)
        if params["precision"] == "bfloat16":
            w = tf1.get_variable('w', [1, nx, nf], initializer=tf.random_normal_initializer(stddev=w_init_stdev, dtype=tf.bfloat16), dtype=tf.bfloat16)
            b = tf1.get_variable('b', [nf], initializer=tf.constant_initializer(0, dtype=tf.bfloat16), dtype=tf.bfloat16)
        else:
            w = tf1.get_variable('w', [1, nx, nf], initializer=tf.random_normal_initializer(stddev=w_init_stdev))
            b = tf1.get_variable('b', [nf], initializer=tf.constant_initializer(0))
        c = tf.reshape(tf.matmul(tf.reshape(x, [-1, nx]), tf.reshape(w, [-1, nf]))+b, start+[nf])
        return c

def attention_mask(nd, ns, *, dtype):
    """1's in the lower triangle, counting from the lower right corner.

    Same as tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd), but doesn't produce garbage on TPUs.
    """
    i = tf.range(nd)[:,None]
    j = tf.range(ns)
    m = i >= j - ns + nd
    return tf.cast(m, dtype)

def attn(x, scope, n_state, *, past, params, batch_size, seq_length, train=False):
    assert x.shape.ndims == 2  # Should be [batch*sequence, features]
    assert n_state % params["n_head"] == 0
    *start, hidden_size = shape_

    num_attention_heads = params["n_head"]
    assert(hidden_size % num_attention_heads == 0)
    size_per_head = hidden_size // num_attention_heads

    if past is not None:
        assert past.shape.ndims == 5  # Should be [batch, 2, heads, sequence, features], where 2 is [k, v]

    def split_heads(x):
        x = tf.reshape(x, [batch_size, seq_length, num_attention_heads, size_per_head])
        x = split_states(x, params["n_head"])
        return tf.transpose(x, [0, 2, 1, 3])

    def merge_heads(x):
        # Reverse of split_heads
        x = tf.transpose(x, [0, 2, 1, 3])
        x = merge_states(x)
        x = tf.reshape(x, [batch_size * seq_length, num_attention_heads * size_per_head])
        return x

    def mask_attn_weights(w):
        # w has shape [batch, heads, dst_sequence, src_sequence], where information flows from src to dst.
        _, _, nd, ns = shape_list(w)
        b = attention_mask(nd, ns, dtype=w.dtype)
        b = tf.reshape(b, [1, 1, nd, ns])
        w = w*b - tf.cast(1e10, w.dtype)*(1-b)
        return w

    def multihead_attn(q, k, v):
        # q, k, v have shape [batch, heads, sequence, features]
        w = tf.matmul(q, k, transpose_b=True)
        w = w * tf.math.rsqrt(tf.cast(shape_list(v)[-1], w.dtype))

        w = mask_attn_weights(w)
        w = softmax(w)

        w = dropout(w, params["attn_dropout"], train)

        a = tf.matmul(w, v)
        return a

    with tf1.variable_scope(scope):
        c = conv1d(x, 'c_attn', n_state*3, params=params)
        q, k, v = map(split_heads, tf.split(c, 3, axis=-1))
        present = tf.stack([k, v], axis=1)
        if past is not None:
            pk, pv = tf.unstack(past, axis=1)
            k = tf.concat([pk, k], axis=-2)
            v = tf.concat([pv, v], axis=-2)
        a = multihead_attn(q, k, v)
        a = merge_heads(a)
        a = conv1d(a, 'c_proj', n_state, params=params)
        a = dropout(a, params["res_dropout"], train)
        return a, present

def mlp(x, scope, n_state, *, params, train=False):
    with tf1.variable_scope(scope):
        nx = shape_list(x)[-1]
        h = gelu(conv1d(x, 'c_fc', n_state, params=params))
        h2 = conv1d(h, 'c_proj', nx, params=params, scale=True)
        h2 = dropout(h2, params["res_dropout"], train)
        return h2

def block(x, scope, *, past, params, attn, train=False, **attn_kws):
    with tf1.variable_scope(scope):
        nx = shape_list(x)[-1]
        ln_1 = norm(x, 'ln_1', params=params)
        a, present = attn(ln_1, 'attn', nx, past=past, params=params, train=train, **attn_kws)
        x = x + a
        ln_2 = norm(x, 'ln_2', params=params)
        m = mlp(ln_2, 'mlp', nx*4, params=params, train=train)
        x = x + m
        return x, present

def past_shape(*, params, batch_size=None, sequence=None):
    return [batch_size, params["n_layer"], 2, params["n_head"], sequence, params["n_embd"] // params["n_head"]]

def expand_tile(value, size):
    """Add a new axis of given size."""
    value = tf.convert_to_tensor(value, name='value')
    ndims = value.shape.ndims
    return tf.tile(tf.expand_dims(value, axis=0), [size] + [1]*ndims)

def positions_for(tokens, past_length):
    batch_size = tf.shape(tokens)[0]
    nsteps = tf.shape(tokens)[1]
    return expand_tile(past_length + tf.range(nsteps), batch_size)

def dropout(x, pdrop, train):
    if train and pdrop > 0:
        x = tf.nn.dropout(x, rate=pdrop)
    return x

def _assert_float_dtype(dtype):
    if not dtype.is_floating:
        raise ValueError("Expected floating point type, got %s." % dtype)
    return dtype

def model(X, params, labels=None, past=None, scope='model', reuse=False, train=False):
    with tf1.variable_scope(scope, reuse=reuse):
        results = {}
        batch, sequence = shape_list(X)

        if params["precision"] == "bfloat16":
            wpe = tf1.get_variable('wpe', [params["n_ctx"], params["n_embd"]], # Position encoding
                             initializer=tf.random_normal_initializer(stddev=0.01, dtype=tf.bfloat16), dtype=tf.bfloat16)
            wte = tf1.get_variable('wte', [params["n_vocab"], params["n_embd"]], # Text encoding
                             initializer=tf.random_normal_initializer(stddev=0.02, dtype=tf.bfloat16), dtype=tf.bfloat16)

        else:
            wpe = tf1.get_variable('wpe', [params["n_ctx"], params["n_embd"]], # Position encoding
                                initializer=tf.random_normal_initializer(stddev=0.01))
            wte = tf1.get_variable('wte', [params["n_vocab"], params["n_embd"]], # Text encoding
                                initializer=tf.random_normal_initializer(stddev=0.02))
        past_length = 0 if past is None else tf.shape(past)[-2]

        wpe = dropout(wpe, params["embed_dropout"], train)
        wte = dropout(wte, params["embed_dropout"], train)

        h = tf.gather(wte, X) + tf.gather(wpe, positions_for(X, past_length))

        ## We keep the representation as a 2D tensor to avoid re-shaping it back and
        ## forth from a 3D tensor to a 2D tensor. Re-shapes are normally free on
        ## the GPU/CPU but may not be free on the TPU, so we want to minimize them to
        ## help the optimizer.
        batch_size, seq_length, hidden_size = shape_list(h)
        h = tf.reshape(h, [batch_size * seq_length, hidden_size])

        # Transformer
        presents = []
        activations = []
        pasts = tf.unstack(past, axis=1) if past is not None else [None] * params["n_layer"]
        assert len(pasts) == params["n_layer"]
        checkpoint=False if 'memory_saving_gradients' not in params else params['memory_saving_gradients']
        every = 1 if 'memory_saving_checkpoints' not in params else params['memory_saving_checkpoints']
        for layer, past in enumerate(pasts):
            def block0(x):
              with tf1.variable_scope(scope, reuse=reuse):
                x1, present = block(x, 'h%d' % layer, past=past, params=params, attn=attn, train=train, batch_size=batch, seq_length=sequence)
                presents.append(present)
                return x1
            @tf.custom_gradient
            def block1(input):
              def grad(dy, variables=None):
                # variables contains the tensors used to calculate
                # first, we use stop_gradient to ensure that the
                input0 = tf.stop_gradient(input)
                # then, we use the disconnected input to recalculate
                # the output for this layer.
                output0 = block0(input0)
                # now that we have the output, we need to calculate
                result = tf.gradients(output0, input0, dy)
                if variables != None:
                  paramcount = lambda vs: sum([np.prod(v.shape.as_list()) for v in vs])
                  def logvars(variables, label, print_variables=True):
                    if print_variables:
                      tf.logging.info("%s (%s parameters): %s", label, paramcount(variables), pps(variables))
                    else:
                      tf.logging.info("%s (%s parameters)", label, paramcount(variables))
                    return variables
                  tf.logging.info("---------")
                  logvars(variables, "block1_grad variables for layer %s" % layer)
                  return result, tf.gradients(output0, variables, dy)
                return result
              output = block0(input)
              return output, grad
            if checkpoint and (isinstance(every, int) and layer % every == 0 or layer in every):
                tf.logging.info('checkpointing layer %d', layer)
                tf.add_to_collection('checkpoints', h)
            if bool(int(os.environ.get('GRADIENT_CHECKPOINTING', '0'))):
              h = block1(h)
            else:
              h = block0(h)
            activations.append(h)
        results['present'] = tf.stack(presents, axis=1)
        results['activations'] = activations
        h = norm(h, 'ln_f', params=params)

        h_flat = tf.reshape(h, [batch*sequence, params["n_embd"]])
        logits = tf.matmul(h_flat, wte, transpose_b=True)
        logits = tf.reshape(logits, [batch, sequence, params["n_vocab"]])
        results['logits'] = logits
        return results


    *start, a, b = shape_list(x)
    return tf.reshape(x, start + [a*b])

import math
import os

import numpy as np
import tensorflow as tf
from pprint import pformat as pps

tf1 = tf.compat.v1

def default_hparams(trainable=True, dtype=tf.float32, scope='model'):
    return {
        'n_vocab': 50257,
        'n_ctx': 1024,
        'n_embd': 768,
        'n_head': 12,
        'n_layer': 12,
        'res_dropout': 0.0,
        'attn_dropout': 0.0,
        'embed_dropout': 0.0,
        'dtype': dtype,
        'trainable': trainable,
        'scope': scope,
        'precision': 'bfloat16' if dtype == tf.bfloat16 else 'float32',
        'scale_by_depth': False,
        'scale_by_in': False,
    }

def shape_list(x):
    """Deal with dynamic shape in tensorflow cleanly."""
    static = x.shape.as_list()
    dynamic = tf.shape(x)
    return [dynamic[i] if s is None else s for i, s in enumerate(static)]

def softmax(x, axis=-1):
    x = x - tf.reduce_max(x, axis=axis, keepdims=True)
    ex = tf.exp(x)
    return ex / tf.reduce_sum(ex, axis=axis, keepdims=True)

def gelu(x):
    return 0.5*x*(1+tf.tanh(np.sqrt(2/np.pi)*(x+0.044715*tf.pow(x, 3))))

def norm(x, scope, *, axis=-1, epsilon=1e-5, params=None):
    """Normalize to mean = 0, std = 1, then do a diagonal affine transform."""
    with tf1.variable_scope(scope):
        n_state = shape_list(x)[-1]
        if params["precision"] == "bfloat16":
            g = tf1.get_variable('g', [n_state], initializer=tf.constant_initializer(1, dtype=tf.bfloat16), dtype=tf.bfloat16)
            b = tf1.get_variable('b', [n_state], initializer=tf.constant_initializer(0, dtype=tf.bfloat16), dtype=tf.bfloat16)
        else:
            g = tf1.get_variable('g', [n_state], initializer=tf.constant_initializer(1))
            b = tf1.get_variable('b', [n_state], initializer=tf.constant_initializer(0))
        u = tf.reduce_mean(x, axis=axis, keepdims=True)
        s = tf.reduce_mean(tf.square(x-u), axis=axis, keepdims=True)
        x = (x - u) * tf.math.rsqrt(s + epsilon)
        x = x*g + b
        return x

def split_states(x, n):
    """Reshape the last dimension of x into [n, x.shape[-1]/n]."""
    *start, u, v = shape_list(x)
    m = u * v
    return tf.reshape(x, start + [n, m//n])

def merge_states(x):
    """Smash the last two dimensions of x into a single dimension."""

def conv1d(x, scope, nf, *, w_init_stdev=0.02, params=None, scale=False):
    if params["scale_by_depth"] and scale: # Scale by sqrt(num_layers), only happens at the final projection before a res block output
        w_init_stdev = w_init_stdev * (1. / math.sqrt(params["n_layer"]))
    if params["scale_by_in"]: # Scale by sqrt(num_input_features)
        w_init_stdev = w_init_stdev * (1. / math.sqrt(shape_list(x)[-1]))

    with tf1.variable_scope(scope):
        *start, nx = shape_list(x)
        if params["precision"] == "bfloat16":
            w = tf1.get_variable('w', [1, nx, nf], initializer=tf.random_normal_initializer(stddev=w_init_stdev, dtype=tf.bfloat16), dtype=tf.bfloat16)
            b = tf1.get_variable('b', [nf], initializer=tf.constant_initializer(0, dtype=tf.bfloat16), dtype=tf.bfloat16)
        else:
            w = tf1.get_variable('w', [1, nx, nf], initializer=tf.random_normal_initializer(stddev=w_init_stdev))
            b = tf1.get_variable('b', [nf], initializer=tf.constant_initializer(0))
        c = tf.reshape(tf.matmul(tf.reshape(x, [-1, nx]), tf.reshape(w, [-1, nf]))+b, start+[nf])
        return c

def attention_mask(nd, ns, *, dtype):
    """1's in the lower triangle, counting from the lower right corner.

    Same as tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd), but doesn't produce garbage on TPUs.
    """
    i = tf.range(nd)[:,None]
    j = tf.range(ns)
    m = i >= j - ns + nd
    return tf.cast(m, dtype)

def attn(x, scope, n_state, *, past, params, batch_size, seq_length, train=False):
    assert x.shape.ndims == 2  # Should be [batch*sequence, features]
    assert n_state % params["n_head"] == 0
    *start, hidden_size = shape_list(x)
    num_attention_heads = params["n_head"]
    assert(hidden_size % num_attention_heads == 0)
    size_per_head = hidden_size // num_attention_heads

    if past is not None:
        assert past.shape.ndims == 5  # Should be [batch, 2, heads, sequence, features], where 2 is [k, v]

    def split_heads(x):
        x = tf.reshape(x, [batch_size, seq_length, num_attention_heads, size_per_head])
        x = split_states(x, params["n_head"])
        return tf.transpose(x, [0, 2, 1, 3])

    def merge_heads(x):
        # Reverse of split_heads
        x = tf.transpose(x, [0, 2, 1, 3])
        x = merge_states(x)
        x = tf.reshape(x, [batch_size * seq_length, num_attention_heads * size_per_head])
        return x

    def mask_attn_weights(w):
        # w has shape [batch, heads, dst_sequence, src_sequence], where information flows from src to dst.
        _, _, nd, ns = shape_list(w)
        b = attention_mask(nd, ns, dtype=w.dtype)
        b = tf.reshape(b, [1, 1, nd, ns])
        w = w*b - tf.cast(1e10, w.dtype)*(1-b)
        return w

    def multihead_attn(q, k, v):
        # q, k, v have shape [batch, heads, sequence, features]
        w = tf.matmul(q, k, transpose_b=True)
        w = w * tf.math.rsqrt(tf.cast(shape_list(v)[-1], w.dtype))

        w = mask_attn_weights(w)
        w = softmax(w)

        w = dropout(w, params["attn_dropout"], train)

        a = tf.matmul(w, v)
        return a

    with tf1.variable_scope(scope):
        c = conv1d(x, 'c_attn', n_state*3, params=params)
        q, k, v = map(split_heads, tf.split(c, 3, axis=-1))
        present = tf.stack([k, v], axis=1)
        if past is not None:
            pk, pv = tf.unstack(past, axis=1)
            k = tf.concat([pk, k], axis=-2)
            v = tf.concat([pv, v], axis=-2)
        a = multihead_attn(q, k, v)
        a = merge_heads(a)
        a = conv1d(a, 'c_proj', n_state, params=params)
        a = dropout(a, params["res_dropout"], train)
        return a, present

def mlp(x, scope, n_state, *, params, train=False):
    with tf1.variable_scope(scope):
        nx = shape_list(x)[-1]
        h = gelu(conv1d(x, 'c_fc', n_state, params=params))
        h2 = conv1d(h, 'c_proj', nx, params=params, scale=True)
        h2 = dropout(h2, params["res_dropout"], train)
        return h2

def block(x, scope, *, past, params, attn, train=False, **attn_kws):
    with tf1.variable_scope(scope):
        nx = shape_list(x)[-1]
        ln_1 = norm(x, 'ln_1', params=params)
        a, present = attn(ln_1, 'attn', nx, past=past, params=params, train=train, **attn_kws)
        x = x + a
        ln_2 = norm(x, 'ln_2', params=params)
        m = mlp(ln_2, 'mlp', nx*4, params=params, train=train)
        x = x + m
        return x, present

def past_shape(*, params, batch_size=None, sequence=None):
    return [batch_size, params["n_layer"], 2, params["n_head"], sequence, params["n_embd"] // params["n_head"]]

def expand_tile(value, size):
    """Add a new axis of given size."""
    value = tf.convert_to_tensor(value, name='value')
    ndims = value.shape.ndims
    return tf.tile(tf.expand_dims(value, axis=0), [size] + [1]*ndims)

def positions_for(tokens, past_length):
    batch_size = tf.shape(tokens)[0]
    nsteps = tf.shape(tokens)[1]
    return expand_tile(past_length + tf.range(nsteps), batch_size)

def dropout(x, pdrop, train):
    if train and pdrop > 0:
        x = tf.nn.dropout(x, rate=pdrop)
    return x

def _assert_float_dtype(dtype):
    if not dtype.is_floating:
        raise ValueError("Expected floating point type, got %s." % dtype)
    return dtype

def model(X, params, labels=None, past=None, scope='model', reuse=False, train=False):
    with tf1.variable_scope(scope, reuse=reuse):
        results = {}
        batch, sequence = shape_list(X)

        if params["precision"] == "bfloat16":
            wpe = tf1.get_variable('wpe', [params["n_ctx"], params["n_embd"]], # Position encoding
                             initializer=tf.random_normal_initializer(stddev=0.01, dtype=tf.bfloat16), dtype=tf.bfloat16)
            wte = tf1.get_variable('wte', [params["n_vocab"], params["n_embd"]], # Text encoding
                             initializer=tf.random_normal_initializer(stddev=0.02, dtype=tf.bfloat16), dtype=tf.bfloat16)

        else:
            wpe = tf1.get_variable('wpe', [params["n_ctx"], params["n_embd"]], # Position encoding
                                initializer=tf.random_normal_initializer(stddev=0.01))
            wte = tf1.get_variable('wte', [params["n_vocab"], params["n_embd"]], # Text encoding
                                initializer=tf.random_normal_initializer(stddev=0.02))
        past_length = 0 if past is None else tf.shape(past)[-2]

        wpe = dropout(wpe, params["embed_dropout"], train)
        wte = dropout(wte, params["embed_dropout"], train)

        h = tf.gather(wte, X) + tf.gather(wpe, positions_for(X, past_length))

        ## We keep the representation as a 2D tensor to avoid re-shaping it back and
        ## forth from a 3D tensor to a 2D tensor. Re-shapes are normally free on
        ## the GPU/CPU but may not be free on the TPU, so we want to minimize them to
        ## help the optimizer.
        batch_size, seq_length, hidden_size = shape_list(h)
        h = tf.reshape(h, [batch_size * seq_length, hidden_size])

        # Transformer
        presents = []
        activations = []
        pasts = tf.unstack(past, axis=1) if past is not None else [None] * params["n_layer"]
        assert len(pasts) == params["n_layer"]
        checkpoint=False if 'memory_saving_gradients' not in params else params['memory_saving_gradients']
        every = 1 if 'memory_saving_checkpoints' not in params else params['memory_saving_checkpoints']
        for layer, past in enumerate(pasts):
            def block0(x):
              with tf1.variable_scope(scope, reuse=reuse):
                x1, present = block(x, 'h%d' % layer, past=past, params=params, attn=attn, train=train, batch_size=batch, seq_length=sequence)
                presents.append(present)
                return x1
            @tf.custom_gradient
            def block1(input):
              def grad(dy, variables=None):
                # variables contains the tensors used to calculate
                # first, we use stop_gradient to ensure that the
                input0 = tf.stop_gradient(input)
                # then, we use the disconnected input to recalculate
                # the output for this layer.
                output0 = block0(input0)
                # now that we have the output, we need to calculate
                result = tf.gradients(output0, input0, dy)
                if variables != None:
                  paramcount = lambda vs: sum([np.prod(v.shape.as_list()) for v in vs])
                  def logvars(variables, label, print_variables=True):
                    if print_variables:
                      tf.logging.info("%s (%s parameters): %s", label, paramcount(variables), pps(variables))
                    else:
                      tf.logging.info("%s (%s parameters)", label, paramcount(variables))
                    return variables
                  tf.logging.info("---------")
                  logvars(variables, "block1_grad variables for layer %s" % layer)
                  return result, tf.gradients(output0, variables, dy)
                return result
              output = block0(input)
              return output, grad
            if checkpoint and (isinstance(every, int) and layer % every == 0 or layer in every):
                tf.logging.info('checkpointing layer %d', layer)
                tf.add_to_collection('checkpoints', h)
            if bool(int(os.environ.get('GRADIENT_CHECKPOINTING', '0'))):
              h = block1(h)
            else:
              h = block0(h)
            activations.append(h)
        results['present'] = tf.stack(presents, axis=1)
        results['activations'] = activations
        h = norm(h, 'ln_f', params=params)

        h_flat = tf.reshape(h, [batch*sequence, params["n_embd"]])
        logits = tf.matmul(h_flat, wte, transpose_b=True)
        logits = tf.reshape(logits, [batch, sequence, params["n_vocab"]])
        results['logits'] = logits
        return results